#!/usr/bin/env python3
import argparse
import os
import re
import subprocess
import sys
import time
import json
from pathlib import Path

VERBOSE = False

def config_path():
    sudo_user = os.environ.get("SUDO_USER")
    home = Path(f"/home/{sudo_user}") if sudo_user else Path.home()
    return home / ".lxcdrc"

def ensure_config():
    path = config_path()
    if path.exists():
        return
    
    config_content = """# lxcd configuration
#
# [workspace] sets the default workspace directory name used by 'lxcd create'
# and 'lxcd enter'.
# Additional sections (e.g. [git], [opencode]) define paths to mount
# into containers via 'lxcd add --map=<section>'.

[workspace]
Code

[git]
.gitconfig

[opencode]
.config/opencode
.local/share/opencode
"""
    try:
        path.write_text(config_content)
        print(f"Created {path}")
    except Exception as e:
        print(f"Warning: Could not create config at {path}: {e}", file=sys.stderr)

def load_config():
    path = config_path()
    if not path.exists():
        return {}
    
    config = {}
    section = None
    values = []
    
    try:
        for line in path.read_text().splitlines():
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            
            m = re.match(r"^\[(.+)\]$", line)
            if m:
                if section is not None:
                    config[section] = values
                section = m.group(1)
                values = []
            elif section is not None:
                values.append(line)
        
        if section is not None:
            config[section] = values
    except Exception as e:
        print(f"Warning: Could not read config from {path}: {e}", file=sys.stderr)
        
    return config

def get_real_user_info():
    """
    Detects the real user when running under sudo.
    Falls back to normal environment variables if not running under sudo.
    """
    sudo_user = os.environ.get("SUDO_USER")
    sudo_uid = os.environ.get("SUDO_UID")
    sudo_gid = os.environ.get("SUDO_GID")

    if sudo_user and sudo_uid and sudo_gid:
        return sudo_user, int(sudo_uid), int(sudo_gid), Path(f"/home/{sudo_user}")
    else:
        return (
            os.environ.get("USER", "root"),
            os.getuid(),
            os.getgid(),
            Path.home()
        )

def detect_timezone():
    if os.environ.get("SNAP"):
        return "UTC"
    
    tz = "UTC"
    tz_file = Path("/etc/timezone")
    if tz_file.exists():
        tz = tz_file.read_text().strip()
    else:
        try:
            result = subprocess.run(
                ["timedatectl"], capture_output=True, text=True, check=False
            )
            m = re.search(r"Time zone:\s+(\S+)", result.stdout)
            if m:
                tz = m.group(1)
        except Exception:
            pass
    return tz

def detect_locale():
    return os.environ.get("LANG", "en_US.UTF-8")

def run_cmd(cmd, check=True):
    """Helper to run shell commands."""
    try:
        return subprocess.run(cmd, check=check, text=True)
    except subprocess.CalledProcessError as e:
        print(f"Error executing command: {' '.join(cmd)}", file=sys.stderr)
        sys.exit(e.returncode)

def lxc_run(args, check=True):
    """Run an lxc command via run_cmd."""
    cmd = ["lxc", *args]
    if VERBOSE:
        print(f"+ {' '.join(cmd)}", file=sys.stderr)
    return run_cmd(cmd, check=check)

def lxc_subprocess(args, **kwargs):
    """Run an lxc command via subprocess.run."""
    cmd = ["lxc", *args]
    if VERBOSE:
        print(f"+ {' '.join(cmd)}", file=sys.stderr)
    return subprocess.run(cmd, **kwargs)

def set_label(name, feature, value="true"):
    """Set a user.lxcd.<feature> config key on a container."""
    lxc_run(["config", "set", name, f"user.lxcd.{feature}", value], check=False)

def container_exists(name):
    result = lxc_subprocess(
        ["info", name],
        capture_output=True, text=True, check=False
    )
    return result.returncode == 0

def handle_list(args):
    print("Fetching lxcd container environments...")

    result = lxc_subprocess(
        ["list", "user.lxcd=true", "--format", "json"],
        capture_output=True, text=True, check=False
    )
    if result.returncode != 0:
        print("Error communicating with LXD.", file=sys.stderr)
        return

    containers = json.loads(result.stdout)

    print(f"\n{'CONTAINER NAME':<25} {'STATUS':<12} {'IPV4 ADDRESS':<25} {'ACTIVE OPTIONS':<25}")
    print("-" * 87)

    count = 0
    cfg = load_config()

    for c in containers:
        name = c["name"]
        status = c["status"]
        ipv4 = "N/A"
        state = c.get("state")
        if state:
            networks = state.get("network") or {}
            eth0 = networks.get("eth0", {})
            addrs = [a["address"] for a in eth0.get("addresses", [])
                     if a.get("family") == "inet"]
            if addrs:
                ipv4 = addrs[0]

        def get_config(key):
            return lxc_subprocess(
                ["config", "get", name, key],
                capture_output=True, text=True, check=False
            ).stdout.strip()

        active_opts = []

        if get_config("security.nesting") == "true":
            active_opts.append("nesting")

        if get_config("user.lxcd.aptcache") == "true":
            active_opts.append("aptcache")

        for section in cfg.keys():
            if section == "workspace":
                continue
            if get_config(f"user.lxcd.map.{section}") == "true":
                active_opts.append(section)

        opts_str = ", ".join(active_opts) if active_opts else "none"
        print(f"{name:<25} {status:<12} {ipv4:<25} {opts_str:<25}")
        count += 1

    if count == 0:
        print("No active lxcd environments found.")
    print()

def handle_clone(args):
    print(f"Cloning container '{args.source}' to '{args.name}'...")

    lxc_run(["copy", args.source, args.name])

    print(f"Starting cloned container '{args.name}'...")
    lxc_run(["start", args.name])
    lxc_run(["exec", args.name, "--", "cloud-init", "status", "--wait"], check=False)

    print(f"Container '{args.name}' cloned successfully from '{args.source}'!")

def handle_create(args):
    ensure_config()
    cfg = load_config()
    workspace = args.workspace or (cfg.get("workspace")[0] if cfg.get("workspace") else "Coder")
    
    host_username, uid, gid, host_home = get_real_user_info()
    if args.minimal:
        args.image = "ubuntu-minimal:"
    host_timezone = detect_timezone()
    host_locale = detect_locale()

    host_coder_dir = host_home / workspace
    container_coder_path = f"/home/{host_username}/{workspace}"

    if not host_coder_dir.exists():
        print(f"Creating missing directory on host: {host_coder_dir}")
        host_coder_dir.mkdir(parents=True, exist_ok=True)
        os.chown(host_coder_dir, uid, gid)

    print(f"Creating LXD container '{args.name}' from image '{args.image}'...")

    apt_proxy_cmd = (
        f"  - mkdir -p /etc/apt/apt.conf.d\n"
        f"  - echo 'Acquire::http::Proxy \"{args.aptcache}\";' > /etc/apt/apt.conf.d/00lxcd-proxy\n"
    ) if args.aptcache else ""

    cloud_config = f"""#cloud-config
timezone: {host_timezone}
locale: {host_locale}
users:
  - name: {host_username}
    gecos: lxcd user
    primary_group: {host_username}
    groups: [adm, sudo, plugdev]
    shell: /bin/bash
    sudo: ALL=(ALL) NOPASSWD:ALL
    lock_passwd: true
    homedir: /home/{host_username}

groups:
  - {host_username}

bootcmd:
  - mkdir -p /home/{host_username}
  - mkdir -p /home/{host_username}/.config /home/{host_username}/.local /home/{host_username}/.local/share
  - chown {uid}:{gid} /home/{host_username}/.config /home/{host_username}/.local /home/{host_username}/.local/share || true
  - chown {uid}:{gid} /home/{host_username}

runcmd:
  - cp -rpn /etc/skel/. /home/{host_username}/
  - mkdir -p {container_coder_path}
  - chown -R {host_username}:{host_username} /home/{host_username}
  - echo "SSH left at default image state"
{apt_proxy_cmd}"""

    lxc_run(["init", args.image, args.name])

    print("Tagging container as managed by lxcd...")
    lxc_run(["config", "set", args.name, "user.lxcd", "true"])

    print(f"Injecting cloud-init config for user '{host_username}'..." )
    lxc_run(["config", "set", args.name, "cloud-init.user-data", cloud_config])

    print(f"Mounting host {host_coder_dir} -> {container_coder_path} (shift=true)...")
    lxc_run(["config", "device", "add", args.name, "host-coder", "disk",
        f"source={host_coder_dir}",
        f"path={container_coder_path}",
        "shift=true"])

    if args.nested:
        print("Enabling nesting capabilities...")
        lxc_run(["config", "set", args.name, "security.nesting", "true"])
        lxc_run(["config", "set", args.name, "security.syscalls.intercept.mknod", "true"])
        set_label(args.name, "nested")

    if args.aptcache:
        print(f"APT cache/proxy will be configured to {args.aptcache} via cloud-init...")
        set_label(args.name, "aptcache")

    print("Starting container and executing cloud-init user provisioning...")
    lxc_run(["start", args.name])
    lxc_run(["exec", args.name, "--", "cloud-init", "status", "--wait"], check=False)

    print(f"Container '{args.name}' is ready! Workspace mounted at {container_coder_path}")

def handle_enter(args):
    cfg = load_config()
    workspace = (cfg.get("workspace")[0] if cfg.get("workspace") else "Coder")
    _, uid, gid, host_home = get_real_user_info()

    state = lxc_subprocess(
        ["info", args.name],
        capture_output=True, text=True, check=False
    )
    if "Status: RUNNING" not in state.stdout:
        print(f"Container '{args.name}' is not running. Starting it...")
        lxc_run(["start", args.name])
        lxc_run(["exec", args.name, "--", "cloud-init", "status", "--wait"], check=False)
        print(f"Container '{args.name}' started.")
    
    try:
        result = lxc_subprocess(
            ["exec", args.name, "--", "id", "-nu", str(uid)],
            capture_output=True, text=True, check=True
        )
        container_username = result.stdout.strip()
    except subprocess.CalledProcessError:
        container_username = "root"

    container_home = f"/home/{container_username}"
    container_coder_dir = f"{container_home}/{workspace}"
    target_cwd = container_coder_dir

    env_prefix = (
        f"HOME={container_home} "
        f"TERM={os.environ.get('TERM', 'xterm-256color')} "
        f"LANG={os.environ.get('LANG', 'C.UTF-8')} "
    )
    
    display = os.environ.get("DISPLAY")
    if display:
        env_prefix += f"DISPLAY={display} "

    shell_cmd = "exec /bin/bash --login"
    exec_cmd = [
        "lxc", "exec", args.name, 
        "--cwd", target_cwd,
        "--user", str(uid),
        "--group", str(gid),
        "--", 
        "sh", "-c", f"env {env_prefix} sh -c '{shell_cmd}'"
    ]
    
    os.execvp("lxc", exec_cmd)

def handle_stop(args):
    print(f"Stopping container '{args.name}'...")
    lxc_run(["stop", args.name], check=False)
    print(f"Container '{args.name}' stopped.")

def handle_start(args):
    print(f"Starting container '{args.name}'...")
    lxc_run(["start", args.name])
    lxc_run(["exec", args.name, "--", "cloud-init", "status", "--wait"], check=False)
    print(f"Container '{args.name}' started.")

def handle_restart(args):
    print(f"Restarting container '{args.name}'...")
    lxc_run(["restart", args.name], check=False)
    lxc_run(["exec", args.name, "--", "cloud-init", "status", "--wait"], check=False)
    print(f"Container '{args.name}' restarted.")

def handle_delete(args):
    print(f"Stopping container '{args.name}'...")
    lxc_run(["stop", args.name], check=False)
    print(f"Deleting container '{args.name}'...")
    lxc_run(["delete", args.name])
    print(f"Container '{args.name}' deleted successfully.")

def handle_add(args):
    host_username, uid, gid, host_home = get_real_user_info()

    if not container_exists(args.name):
        print(f"Error: Container '{args.name}' does not exist.", file=sys.stderr)
        sys.exit(1)

    if not args.nested and not args.aptcache and not args.map and not args.unmap:
        print("Error: Specify at least one option to add.\n"
              "Example: lxcd add <name> --nested --aptcache", file=sys.stderr)
        sys.exit(1)

    state = lxc_subprocess(
        ["info", args.name],
        capture_output=True, text=True, check=False
    )
    was_stopped = "Status: RUNNING" not in state.stdout
    if was_stopped:
        print(f"Starting container '{args.name}' temporarily to apply system operations...")
        lxc_subprocess(["start", args.name], check=False)
        time.sleep(2)

    if args.nested:
        print(f"Enabling nesting support for '{args.name}'...")
        lxc_run(["config", "set", args.name, "security.nesting", "true"])
        lxc_run(["config", "set", args.name, "security.syscalls.intercept.mknod", "true"])
        set_label(args.name, "nested")

    if args.aptcache:
        print(f"Configuring APT proxy to {args.aptcache} inside '{args.name}'...")
        lxc_run(["exec", args.name, "--", "mkdir", "-p", "/etc/apt/apt.conf.d"])
        lxc_run(["exec", args.name, "--", "sh", "-c",
                 f"echo 'Acquire::http::Proxy \"{args.aptcache}\";' > /etc/apt/apt.conf.d/00lxcd-proxy"])
        set_label(args.name, "aptcache")

    if args.map:
        cfg = load_config()
        paths = cfg.get(args.map)
        if paths:
            for idx, path in enumerate(paths):
                host_path = host_home / path
                if host_path.exists():
                    container_path = f"/home/{host_username}/{path}"
                    lxc_subprocess(["config", "device", "remove", args.name, f"map-{args.map}-{idx}"],
                                   capture_output=True, text=True, check=False)
                    lxc_run(["config", "device", "add", args.name, f"map-{args.map}-{idx}", "disk",
                             f"source={host_path}",
                             f"path={container_path}",
                             "shift=true"])
                    print(f"Mapped {host_path} -> {container_path}")
                else:
                    print(f"Warning: {host_path} not found, skipping", file=sys.stderr)
            set_label(args.name, f"map.{args.map}")
        else:
            print(f"Warning: Section '{args.map}' not found in ~/.lxcdrc", file=sys.stderr)

    if args.unmap:
        idx = 0
        while True:
            result = lxc_subprocess(["config", "device", "remove", args.name, f"map-{args.unmap}-{idx}"],
                                    capture_output=True, text=True, check=False)
            if result.returncode != 0:
                break
            idx += 1
        set_label(args.name, f"map.{args.unmap}", "")
        print(f"Unmapped section '{args.unmap}'")

    if was_stopped:
        print(f"Returning '{args.name}' to its original stopped state...")
        lxc_subprocess(["stop", args.name], check=False)

    print(f"\nSuccessfully updated container '{args.name}' with requested features!")

def handle_rename(args):
    print(f"Preparing to rename '{args.old_name}' to '{args.new_name}'...")

    print(f"Stopping '{args.old_name}'...")
    lxc_subprocess(
        ["stop", args.old_name],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL
    )

    lxc_run(["rename", args.old_name, args.new_name])

    print(f"Starting renamed container '{args.new_name}'...")
    lxc_run(["start", args.new_name])
    lxc_run(["exec", args.new_name, "--", "cloud-init", "status", "--wait"], check=False)

    print(f"Container successfully renamed to '{args.new_name}'!")

def main():
    parser = argparse.ArgumentParser(description="lxcd: A Distrobox-like wrapper for LXD containers")
    parser.add_argument("-v", "--verbose", action="store_true", help="Show verbose command output")
    subparsers = parser.add_subparsers(dest="command", required=True)

    # Create command
    parser_create = subparsers.add_parser("create", help="Create a new LXD container")
    parser_create.add_argument("name", help="Name of the container")
    parser_create.add_argument("-i", "--image", default="ubuntu:", help="LXD image to use (defaults to 'ubuntu:')")
    parser_create.add_argument("-m", "--minimal", action="store_true", help="Use ubuntu-minimal: image instead")
    parser_create.add_argument("-w", "--workspace", help="Workspace directory name (default: Coder or ~/.lxcdrc [workspace])")
    parser_create.add_argument("--nested", action="store_true", help="Enable nesting")
    parser_create.add_argument("--aptcache", help="URL for an APT cache/proxy (e.g. http://10.0.0.1:3142)")

    # Clone command
    parser_clone = subparsers.add_parser("clone", help="Clone an existing LXD container")
    parser_clone.add_argument("source", help="Name of the source container")
    parser_clone.add_argument("name", help="Name of the new cloned container")

    # Rename command
    parser_rename = subparsers.add_parser("rename", aliases=["mv"], help="Rename an existing LXD container")
    parser_rename.add_argument("old_name", help="Current name of the container")
    parser_rename.add_argument("new_name", help="New name for the container")

    # Enter command
    parser_enter = subparsers.add_parser("enter", help="Enter an existing LXD container")
    parser_enter.add_argument("name", help="Name of the container to enter")

    # List command
    parser_list = subparsers.add_parser("list", aliases=["ls"], help="List all managed LXD containers")

    # Stop command
    parser_stop = subparsers.add_parser("stop", help="Stop a running container")
    parser_stop.add_argument("name", help="Name of the container to stop")

    # Start command
    parser_start = subparsers.add_parser("start", help="Start a stopped container")
    parser_start.add_argument("name", help="Name of the container to start")

    # Restart command
    parser_restart = subparsers.add_parser("restart", help="Restart a container")
    parser_restart.add_argument("name", help="Name of the container to restart")

    # Delete command
    parser_delete = subparsers.add_parser("delete", help="Delete an LXD container")
    parser_delete.add_argument("name", help="Name of the container to delete")

    # Add command
    parser_add = subparsers.add_parser("add", help="Add features to an existing container")
    parser_add.add_argument("name", help="Name of the container")
    parser_add.add_argument("--nested", action="store_true", help="Enable nesting")
    parser_add.add_argument("--aptcache", help="URL for an APT cache/proxy (e.g. http://10.0.0.1:3142)")
    parser_add.add_argument("--map", help="Map paths from a ~/.lxcdrc section into the container")
    parser_add.add_argument("--unmap", help="Remove a previously mapped section")

    args = parser.parse_args()

    global VERBOSE
    VERBOSE = args.verbose

    if args.command == "create":
        handle_create(args)
    elif args.command == "clone":
        handle_clone(args)
    elif args.command in ["rename", "mv"]:
        handle_rename(args)
    elif args.command == "enter":
        handle_enter(args)
    elif args.command in ["list", "ls"]:
        handle_list(args)
    elif args.command == "stop":
        handle_stop(args)
    elif args.command == "start":
        handle_start(args)
    elif args.command == "restart":
        handle_restart(args)
    elif args.command == "delete":
        handle_delete(args)
    elif args.command == "add":
        handle_add(args)

if __name__ == "__main__":
    main()