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

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():
    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 FileNotFoundError:
            pass
    return tz


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


VERBOSE = False

def lxc_prefix():
    """Return ['sudo'] if the user isn't root and isn't in the lxd group."""
    if os.getuid() == 0:
        return []
    try:
        if grp.getgrnam("lxd").gr_gid in os.getgroups():
            return []
    except KeyError:
        pass
    return ["sudo"]

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, auto-adding sudo if needed."""
    cmd = [*lxc_prefix(), "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, auto-adding sudo if needed."""
    cmd = [*lxc_prefix(), "lxc", *args]
    if VERBOSE:
        print(f"+ {' '.join(cmd)}", file=sys.stderr)
    return subprocess.run(cmd, **kwargs)

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

    import json
    containers = json.loads(result.stdout)

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

    count = 0
    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]

        # Read lxcd feature labels via individual config get
        def get_config(key):
            return lxc_subprocess(
                ["config", "get", name, key],
                capture_output=True, text=True, check=False
            ).stdout.strip()

        active_opts = []

        # Check nesting via security.nesting (native LXD config)
        if get_config("security.nesting") == "true":
            active_opts.append("nesting")

        # Check labels for other features
        for feature in ("certs", "git", "aptcache", "ssh-client"):
            if get_config(f"user.lxcd.{feature}") == "true":
                active_opts.append(feature)

        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}'...")

    # 1. Native LXD copy (duplicates the filesystem, mounts, and tags)
    lxc_run([ "copy", args.source, args.name])

    # 2. Start the new cloned container
    print(f"Starting cloned container '{args.name}'...")
    lxc_run([ "start", args.name])

    # 3. Wait for network/cloud-init stabilization
    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):
    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()

    # Paths for targeted workspace sharing
    host_coder_dir = host_home / "Coder"
    container_coder_path = f"/home/{host_username}/Coder"

    # Ensure the 'Coder' directory exists on the host machine
    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}'...")

    # Cloud-init configuration
    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 ""

    ssh_client_packages = "packages:\n  - openssh-client\n" if args.ssh_client else ""
    ssh_client_cmd = (
        f"  - mkdir -p /home/{host_username}/.ssh\n"
        f"  - chmod 700 /home/{host_username}/.ssh\n"
    ) if args.ssh_client 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}

{ssh_client_packages}bootcmd:
  - mkdir -p /home/{host_username}
  - chown {uid}:{gid} /home/{host_username}

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

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

    # Tag the container so lxcd knows it owns it
    print("Tagging container as managed by lxcd...")
    lxc_run([ "config", "set", args.name, "user.lxcd", "true"])

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

    # Mount host's Coder directory using shift=true
    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"])

    # Handle --git option
    if args.git:
        gitconfig_path = host_home / ".gitconfig"
        if gitconfig_path.exists():
            container_git_path = f"/home/{host_username}/.gitconfig"
            print(f"Mapping host git configurations to {container_git_path} (shift=true)...")
            lxc_run(["config", "device", "add", args.name, "gitconfig", "disk",
                f"source={gitconfig_path}",
                f"path={container_git_path}",
                "shift=true"])
            set_label(args.name, "git")
        else:
            print(f"Warning: --git requested, but {gitconfig_path} was not found on host.", file=sys.stderr)

    # Handle --nested option
    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")

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

    # Handle --ssh-client option
    if args.ssh_client:
        print(f"SSH client and ~/.ssh directory will be set up for '{host_username}' via cloud-init...")
        set_label(args.name, "ssh-client")

    # Handle --certs option
    if args.certs:
        host_certs = Path("/etc/ssl/certs/ca-certificates.crt")
        if host_certs.exists():
            print("Mounting host ca-certificates bundle securely into container...")
            lxc_run(["config", "device", "add", args.name, "host-certs", "disk",
                f"source={host_certs}",
                f"path={host_certs}",
                "shift=true"])
            set_label(args.name, "certs")
        else:
            print("Warning: Host cert bundle path not found. Skipping certificate mount configuration.")

    # Start container and await cloud-init provisioning
    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):
    _, uid, gid, host_home = get_real_user_info()
    host_cwd = os.getcwd()
    
    # Dynamically find the user matching your UID inside the container
    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}"

    # Smart CWD Mapping: Ensure the host directory actually exists in the container
    host_coder_dir = str(host_home / "Coder")
    container_coder_dir = f"{container_home}/Coder"
    
    # If the user is inside the mapped Coder directory on the host, translate the path
    if host_cwd.startswith(host_coder_dir):
        target_cwd = host_cwd.replace(host_coder_dir, container_coder_dir, 1)
    else:
        # Fallback to the container's home directory if executed from an unmapped host path
        target_cwd = container_home

    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} "

    # Check if container has the ssh-client feature (set by --ssh-client)
    has_ssh_client = lxc_subprocess(
        ["config", "get", args.name, "user.lxcd.ssh-client"],
        capture_output=True, text=True, check=False
    ).stdout.strip() == "true"

    # Execute using native LXD flags for user, group, and working directory
    prefix = lxc_prefix()
    shell_cmd = "exec /bin/bash --login"
    if has_ssh_client:
        shell_cmd = "eval $(ssh-agent) && ssh-add; " + shell_cmd
    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}'"
    ]
    
    if prefix:
        os.execvp("sudo", ["sudo", *exec_cmd])
    else:
        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 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_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.git and not args.certs and not args.aptcache and not args.ssh_client:
        print("Error: Specify at least one option to add.\n"
              "Example: lxcd add <name> --git --nested", 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.git:
        was_running = not was_stopped
        host_gitconfig = host_home / ".gitconfig"
        if host_gitconfig.exists():
            print(f"Mapping host .gitconfig into '{args.name}' as a device mount...")
            lxc_subprocess(["stop", args.name], check=False)
            lxc_subprocess(["config", "device", "remove", args.name, "gitconfig"],
                           capture_output=True, text=True, check=False)
            lxc_run(["config", "device", "add", args.name, "gitconfig", "disk",
                     f"source={host_gitconfig}",
                     f"path=/home/{host_username}/.gitconfig",
                     "shift=true"])
            set_label(args.name, "git")
            if was_running:
                lxc_run([ "start", args.name])
                lxc_run([ "exec", args.name, "--", "cloud-init", "status", "--wait"], check=False)
        else:
            print(f"Warning: Host ~/.gitconfig not found. Skipping git config mapping.")

    if args.certs:
        host_certs = Path("/etc/ssl/certs/ca-certificates.crt")
        if host_certs.exists():
            print(f"Mounting host ca-certificates bundle into '{args.name}'...")
            lxc_subprocess(["config", "device", "remove", args.name, "host-certs"],
                           capture_output=True, text=True, check=False)
            lxc_run(["config", "device", "add", args.name, "host-certs", "disk",
                f"source={host_certs}",
                f"path={host_certs}",
                "shift=true"])
            set_label(args.name, "certs")
        else:
            print("Warning: Host cert bundle path not found. Skipping certificate configuration.")

    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.ssh_client:
        print(f"Installing SSH client and setting up ~/.ssh inside '{args.name}'...")
        lxc_run([ "exec", args.name, "--", "env", "DEBIAN_FRONTEND=noninteractive", "apt-get", "update", "-y"])
        lxc_run([ "exec", args.name, "--", "env", "DEBIAN_FRONTEND=noninteractive", "apt-get", "install", "-y", "openssh-client"])
        lxc_run([ "exec", args.name, "--", "mkdir", "-p", f"/home/{host_username}/.ssh"])
        lxc_run([ "exec", args.name, "--", "chmod", "700", f"/home/{host_username}/.ssh"])
        lxc_run([ "exec", args.name, "--", "chown", f"{host_username}:{host_username}", f"/home/{host_username}/.ssh"])
        set_label(args.name, "ssh-client")

    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}'...")

    # 1. LXD requires containers to be stopped before renaming
    print(f"Stopping '{args.old_name}'...")
    # Using stdout/stderr redirection internally to hide the error if it's already stopped
    lxc_subprocess(
        ["stop", args.old_name],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL
    )

    # 2. Execute the rename
    lxc_run([ "rename", args.old_name, args.new_name])

    # 3. Restart the container under its new name
    print(f"Starting renamed container '{args.new_name}'...")
    lxc_run([ "start", args.new_name])

    # 4. Wait for it to stabilize
    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 (Sudo Optimized)")
    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("--nested", action="store_true", help="Enable nesting")
    parser_create.add_argument("--git", action="store_true", help="Map host ~/.gitconfig into the container")
    parser_create.add_argument("--certs", action="store_true", help="Mount host CA certificates bundle")
    parser_create.add_argument("--aptcache", help="URL for an APT cache/proxy (e.g. http://10.0.0.1:3142)")
    parser_create.add_argument("--ssh-client", action="store_true",
        help="Install SSH client and set up ~/.ssh for remote access from within the container")

    # 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("--git", action="store_true", help="Map host ~/.gitconfig into the container")
    parser_add.add_argument("--certs", action="store_true", help="Mount host CA certificates bundle")
    parser_add.add_argument("--aptcache", help="URL for an APT cache/proxy (e.g. http://10.0.0.1:3142)")
    parser_add.add_argument("--ssh-client", action="store_true",
        help="Install SSH client and set up ~/.ssh for remote access from within the container")

    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()