#!/usr/bin/env python3
import argparse
import os
import re
import subprocess
import sys
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(force=False, workspace="Coder", packages=None, nested=False,
                  image=None, aptcache=None, ssh=None):
    path = config_path()
    if path.exists() and not force:
        return
    
    packages = packages or []
    packages_block = "\n".join(packages) if packages else ""
    nested_value = "true" if nested else ""
    image_value = image or "ubuntu:"
    aptcache_value = aptcache or ""
    ssh_value = ssh or ""

    config_content = f"""# lxcd configuration
#
# [options] holds key=value defaults applied by 'lxcd create'.
#   workspace: default workspace directory name
#   image:     LXD image to use
#   nested:    enable nesting (e.g. true) or leave blank
#   aptcache:  APT cache/proxy URL or leave blank
#   ssh:       SSH public key to authorize for the container user; blank
#              disables the SSH server
#
# [packages] lists packages installed by default in new containers via
# 'lxcd create'.
#
# [mappings] lists paths (relative to ~/) mounted into containers when they
# are entered via 'lxcd enter'.

[options]
workspace={workspace}
image={image_value}
nested={nested_value}
aptcache={aptcache_value}
ssh={ssh_value}

[packages]
{packages_block}

[mappings]
"""
    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 require_config():
    path = config_path()
    if not path.exists():
        print(f"Error: Config file {path} not found. Run 'lxcd init' to create it.", file=sys.stderr)
        sys.exit(1)

def load_config():
    path = config_path()
    if not path.exists():
        return {}
    
    config = {}
    section = None
    entries = []

    def flush():
        nonlocal entries
        if section is None:
            return
        kv = {}
        is_dict = False
        for e in entries:
            if "=" in e:
                key, _, value = e.partition("=")
                kv[key.strip()] = value.strip()
                is_dict = True
        if is_dict:
            config[section] = kv
        else:
            config[section] = list(entries)
        entries = []

    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:
                flush()
                section = m.group(1)
                entries = []
            elif section is not None:
                entries.append(line)

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

    COLUMNS = {
        "name":     ("CONTAINER NAME", 25),
        "status":   ("STATUS", 12),
        "ipv4":     ("IPV4 ADDRESS", 22),
        "image":    ("IMAGE", 18),
        "packages": ("PACKAGES", 25),
        "opts":     ("ACTIVE OPTIONS", 25),
    }

    if args.brief is not None:
        cols = [c.strip() for c in args.brief.split(",")] if args.brief.strip() else ["ipv4"]
    else:
        cols = ["name", "status", "ipv4", "image", "packages", "opts"]

    if "name" not in cols:
        cols = ["name", *cols]

    seen = set()
    cols = [c for c in cols if not (c in seen or seen.add(c))]

    chosen = {c: COLUMNS[c] for c in cols if c in COLUMNS}

    header = " ".join(f"{label:<{w}}" for label, w in chosen.values())
    print(f"\n{header}")
    print("-" * sum(w for _, w in chosen.values()))

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

    count = 0
    cfg = load_config()
    map_sections = [k for k in cfg.keys() if k not in ("options", "packages")]

    for container in containers:
        name = container["name"]
        status = container["status"]
        ipv4 = "N/A"
        container_state = container.get("state")
        if container_state:
            networks = container_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]

        cells = {
            "name": name,
            "status": status,
            "ipv4": ipv4,
        }

        if "opts" in cols:
            active = []
            if get_config(name, "security.nesting") == "true":
                active.append("nesting")
            if get_config(name, "user.lxcd.aptcache") == "true":
                active.append("aptcache")
            if get_config(name, "user.lxcd.ssh") == "true":
                active.append("ssh")
            for section in map_sections:
                if get_config(name, f"user.lxcd.map.{section}") == "true":
                    active.append(section)
            cells["opts"] = ", ".join(active) if active else "none"

        if "image" in cols:
            cells["image"] = get_config(name, "user.lxcd.image") or "N/A"

        if "packages" in cols:
            cells["packages"] = get_config(name, "user.lxcd.packages") or "none"

        row = " ".join(f"{cells[c]:<{w}}" for c, (label, w) in chosen.items())
        print(row)
        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_init(args):
    if args.edit:
        path = config_path()
        if not path.exists():
            packages = [p.strip() for p in args.packages.split(",")] if args.packages else None
            ensure_config(force=True, workspace=args.workspace, packages=packages,
                          nested=args.nested, image=args.image,
                          aptcache=args.aptcache, ssh=args.ssh)
        editor = os.environ.get("EDITOR", os.environ.get("VISUAL", "vi"))
        subprocess.run([editor, str(path)])
        return
    if config_path().exists():
        if args.force:
            print(f"Overwriting existing config at {config_path()}...")
        else:
            print(f"Error: Config file {config_path()} already exists. Use --force to overwrite.", file=sys.stderr)
            sys.exit(1)
    packages = [p.strip() for p in args.packages.split(",")] if args.packages else None
    ensure_config(force=args.force, workspace=args.workspace, packages=packages,
                  nested=args.nested, image=args.image,
                  aptcache=args.aptcache, ssh=args.ssh)

def handle_create(args):
    cfg = load_config()
    options = cfg.get("options") if isinstance(cfg.get("options"), dict) else {}
    workspace = args.workspace or options.get("workspace") or "Coder"
    nested = args.nested or bool(options.get("nested"))
    aptcache = args.aptcache or options.get("aptcache") or None
    image = args.image or options.get("image") or "ubuntu:"
    ssh_key = args.ssh or options.get("ssh") or None
    
    host_username, uid, gid, host_home = get_real_user_info()
    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 '{image}'...")

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

    if args.packages:
        packages = [p.strip() for p in args.packages.split(",") if p.strip()]
    else:
        packages = cfg.get("packages") or []
    if ssh_key and not any("openssh-server" in p for p in packages):
        packages = [*packages, "openssh-server"]

    packages_cmd = ""
    if packages:
        pkg_list = "\n".join(f"  - {pkg}" for pkg in packages)
        packages_cmd = f"packages:\n{pkg_list}\n"

    if ssh_key:
        ssh_auth_cmd = f"    ssh_authorized_keys:\n      - {ssh_key}\n"
        ssh_runcmd = "  - systemctl enable --now sshd || systemctl enable --now ssh || true"
        print("Enabling SSH server and authorizing provided public key...")
    else:
        ssh_auth_cmd = ""
        ssh_runcmd = "  - systemctl disable --now sshd 2>/dev/null || systemctl disable --now ssh 2>/dev/null || true"

    cloud_config = f"""#cloud-config
timezone: {host_timezone}
locale: {host_locale}
{packages_cmd}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}
{ssh_auth_cmd}
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}
{ssh_runcmd}
{apt_proxy_cmd}"""

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

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

    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 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 aptcache:
        print(f"APT cache/proxy will be configured to {aptcache} via cloud-init...")
        set_label(args.name, "aptcache")

    if ssh_key:
        set_label(args.name, "ssh")

    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 apply_mappings(name, host_username, host_home):
    cfg = load_config()
    paths = cfg.get("mappings")
    if not isinstance(paths, list) or not paths:
        return
    for idx, spec in enumerate(paths):
        is_dir = spec.endswith("/")
        if spec.startswith("/"):
            host_path = spec.rstrip("/")
            container_path = spec
        else:
            host_path = str(host_home / spec.rstrip("/"))
            container_path = f"/home/{host_username}/{spec}"
        lxc_subprocess(["config", "device", "remove", name, f"map-{idx}"],
                       capture_output=True, text=True, check=False)
        if os.path.exists(host_path):
            lxc_run(["config", "device", "add", name, f"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(name, "map.mappings")

def handle_enter(args):
    cfg = load_config()
    options = cfg.get("options") if isinstance(cfg.get("options"), dict) else {}
    workspace = options.get("workspace") or "Coder"
    host_username, 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.")

    apply_mappings(args.name, host_username, host_home)
    
    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} "

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

def resolve_targets(args):
    try:
        names = list(args.names)
    except AttributeError:
        names = [args.name]
    if args.all:
        result = lxc_subprocess(
            ["list", "user.lxcd=true", "-c", "n", "--format", "csv"],
            capture_output=True, text=True, check=False
        )
        if result.returncode != 0:
            print("Error communicating with LXD.", file=sys.stderr)
            sys.exit(1)
        names += [n.strip() for n in result.stdout.splitlines() if n.strip()]

    seen = set()
    unique = []
    for n in names:
        if n not in seen:
            seen.add(n)
            unique.append(n)

    if not unique:
        print("Error: Specify at least one container name or use --all.", file=sys.stderr)
        sys.exit(1)
    return unique

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

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

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

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

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():
    description = """lxcd: A Distrobox-like wrapper for LXD containers.

This tool simplifies managing LXD-based development environments.
It supports configuration-driven workspace names and dynamic path mapping
controlled by a ~/.lxcdrc configuration file.

Configuration file (~/.lxcdrc) format:
  [options]
  workspace=Coder
  nested=
  aptcache=

  [packages]
  vim
  git

  [mappings]
  .gitconfig

Usage Workflow:
  1. Create a container:  lxcd create myenv
  2. Enter container:     lxcd enter myenv
     (paths in [mappings] are mounted automatically on enter)"""

    parser = argparse.ArgumentParser(
        description=description,
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("-v", "--verbose", action="store_true", help="Show verbose command output")
    subparsers = parser.add_subparsers(dest="command", required=True)

    # Init command
    parser_init = subparsers.add_parser("init", help="Write the ~/.lxcdrc config file")
    parser_init.add_argument("-f", "--force", action="store_true", help="Overwrite an existing config file")
    parser_init.add_argument("-w", "--workspace", default="Coder", help="Default workspace directory name (default: Coder)")
    parser_init.add_argument("-e", "--edit", action="store_true", help="Edit the config file with $EDITOR (creates it first if missing)")
    parser_init.add_argument("-p", "--packages", help="Comma-separated list of packages installed by default in new containers")
    parser_init.add_argument("-n", "--nested", action="store_true", help="Enable nesting by default in new containers")
    parser_init.add_argument("-i", "--image", help="LXD image to use by default in new containers (default: ubuntu:)")
    parser_init.add_argument("--aptcache", help="URL for an APT cache/proxy stored as a default (e.g. http://10.0.0.1:3142)")
    parser_init.add_argument("--ssh", help="SSH public key authorized for the container user by default")

    # 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", help="LXD image to use (default: ubuntu: or ~/.lxcdrc [options] image=)")
    parser_create.add_argument("-w", "--workspace", help="Workspace directory name (default: ~/.lxcdrc [options] 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)")
    parser_create.add_argument("-p", "--packages", help="Comma-separated packages to install (overrides ~/.lxcdrc [packages])")
    parser_create.add_argument("--ssh", help="SSH public key to authorize (default: ~/.lxcdrc [options] ssh=; blank disables SSH)")

    # 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")
    parser_list.add_argument("-b", "--brief", nargs="?", const="", metavar="COLS",
                             help="Only show name plus a comma-separated subset of status,ipv4,image,packages,opts (default: ipv4)")

    # Stop command
    parser_stop = subparsers.add_parser("stop", help="Stop one or more running containers")
    parser_stop.add_argument("names", nargs="*", help="Names of containers to stop")
    parser_stop.add_argument("-a", "--all", action="store_true", help="Stop all lxcd-managed containers")

    # Start command
    parser_start = subparsers.add_parser("start", help="Start one or more stopped containers")
    parser_start.add_argument("names", nargs="*", help="Names of containers to start")
    parser_start.add_argument("-a", "--all", action="store_true", help="Start all lxcd-managed containers")

    # Restart command
    parser_restart = subparsers.add_parser("restart", help="Restart one or more containers")
    parser_restart.add_argument("names", nargs="*", help="Names of containers to restart")
    parser_restart.add_argument("-a", "--all", action="store_true", help="Restart all lxcd-managed containers")

    # Delete command
    parser_delete = subparsers.add_parser("delete", help="Delete one or more LXD containers")
    parser_delete.add_argument("names", nargs="*", help="Names of containers to delete")
    parser_delete.add_argument("-a", "--all", action="store_true", help="Delete all lxcd-managed containers")

    args = parser.parse_args()

    global VERBOSE
    VERBOSE = args.verbose

    if args.command != "init":
        require_config()

    if args.command == "init":
        handle_init(args)
    elif 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)

if __name__ == "__main__":
    main()