#!/usr/bin/env python3
"""Adopt existing systemd-nspawn containers into the gaol store.

Reads .nspawn config files from /etc/systemd/nspawn/, extracts distro info
from each container's /etc/os-release, and retrieves creation timestamps
from machinectl. Writes gaol store YAMLs and .gaol.json markers.

Must be run as root (needs access to /var/lib/machines/).
"""

from __future__ import annotations

import argparse
import contextlib
import json
import os
import re
import subprocess
import sys
from datetime import datetime
from pathlib import Path

import yaml

NSPAWN_DIR = Path("/etc/systemd/nspawn")
MACHINES_DIR = Path("/var/lib/machines")
DEFAULT_SKIP = {"alertmanager", "dx-template"}


def get_store_dir() -> Path:
    """Get the gaol store directory for the real (non-root) user."""
    sudo_user = os.environ.get("SUDO_USER")
    home = Path(f"/home/{sudo_user}") if sudo_user else Path.home()
    return home / ".config" / "gaol" / "containers"


def parse_nspawn_config(path: Path) -> dict:
    """Parse a .nspawn config file, handling duplicate keys like Bind= and Environment=."""
    config: dict = {
        "exec": {},
        "network": {},
        "files": {"binds": [], "bind_ro": []},
    }

    current_section: str | None = None
    environments: list[tuple[str, str]] = []
    capabilities: list[str] = []

    for line in path.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue

        if line.startswith("[") and line.endswith("]"):
            current_section = line[1:-1].lower()
            continue

        if "=" not in line:
            continue

        key, _, value = line.partition("=")
        key = key.strip()
        value = value.strip()

        if current_section == "exec":
            if key == "Environment":
                env_key, _, env_val = value.partition("=")
                environments.append((env_key.strip(), env_val.strip()))
            elif key == "Capability":
                capabilities.extend(value.split())
            else:
                config["exec"][key] = value

        elif current_section == "network":
            if key == "Port":
                # Port=tcp:host:container
                parts = value.split(":")
                if len(parts) == 3:
                    host_port = int(parts[1])
                    container_port = int(parts[2])
                    config["network"].setdefault("ports", {})[host_port] = container_port
            else:
                config["network"][key] = value

        elif current_section == "files" and key in ("Bind", "BindReadOnly"):
            # Bind=/host:/container or Bind=/path (same path in container)
            parts = value.split(":")
            if len(parts) >= 2:
                host_path = parts[0]
                container_path = ":".join(parts[1:])
            else:
                host_path = parts[0]
                container_path = parts[0]

            target = "bind_ro" if key == "BindReadOnly" else "binds"
            config["files"][target].append((host_path, container_path))

    config["exec"]["_environments"] = environments
    config["exec"]["_capabilities"] = capabilities
    return config


def determine_network_mode(network: dict) -> str:
    """Map nspawn network settings to gaol NetworkMode value.

    Priority:
      1. Private=no → host
      2. Bridge=<name> → bridged
      3. VirtualEthernet=yes (no bridge) → shared
      4. Private=yes (no veth) → none
      5. VirtualEthernet=no → host
      6. No config → shared (gaol default)
    """
    private = network.get("Private", "").lower()
    veth = network.get("VirtualEthernet", "").lower()
    bridge = network.get("Bridge", "")

    if private == "no":
        return "host"
    if bridge:
        return "bridged"
    if veth == "yes":
        return "shared"
    if private == "yes":
        return "none"
    if veth == "no":
        return "host"
    return "shared"


def read_distro(name: str) -> str:
    """Read VERSION_CODENAME from the container's /etc/os-release."""
    os_release = MACHINES_DIR / name / "etc" / "os-release"
    if not os_release.exists():
        return "trixie"
    for line in os_release.read_text().splitlines():
        if line.startswith("VERSION_CODENAME="):
            return line.split("=", 1)[1].strip().strip('"')
    return "trixie"


def get_image_timestamps() -> dict[str, datetime]:
    """Parse machinectl list-images for creation timestamps."""
    result = subprocess.run(
        ["machinectl", "list-images", "--no-legend"],
        capture_output=True,
        text=True,
        check=True,
    )
    timestamps: dict[str, datetime] = {}
    for line in result.stdout.splitlines():
        parts = line.split()
        if len(parts) < 6:
            continue
        name = parts[0]
        # Format: name type ro dow YYYY-MM-DD HH:MM:SS TZ size
        for i, p in enumerate(parts):
            if re.match(r"\d{4}-\d{2}-\d{2}", p) and i + 1 < len(parts):
                with contextlib.suppress(ValueError):
                    timestamps[name] = datetime.strptime(f"{p} {parts[i + 1]}", "%Y-%m-%d %H:%M:%S")
                break
    return timestamps


def detect_gpu(binds: list[tuple[str, str]]) -> dict | None:
    """Detect GPU config from bind mounts."""
    has_kfd = any(h == "/dev/kfd" for h, _ in binds)
    if has_kfd:
        return {
            "enabled": True,
            "vendor": "amd",
            "device_ids": [],
            "capabilities": ["compute", "utility"],
            "require_nvidia_runtime": True,
            "memory_limit": None,
        }
    # TODO: NVIDIA detection via /dev/nvidia* binds
    return None


def build_store_yaml(
    name: str,
    config: dict,
    distro: str,
    created_at: datetime,
    now: datetime,
) -> dict:
    """Build store YAML matching StoredContainer.model_dump(mode='json')."""
    network = config["network"]
    exec_cfg = config["exec"]
    files = config["files"]

    mode = determine_network_mode(network)
    # model_dump(mode="json") converts int dict keys to strings
    ports = {str(k): v for k, v in network.get("ports", {}).items()}

    volumes: dict[str, str] = {}
    for host_path, container_path in files["binds"] + files["bind_ro"]:
        volumes[host_path] = container_path

    environment: dict[str, str] = {}
    for key, val in exec_cfg.get("_environments", []):
        environment[key] = val

    private_users_raw = exec_cfg.get("PrivateUsers", "").lower()
    if private_users_raw == "no":
        private_users = False
    elif private_users_raw == "yes":
        private_users = True
    else:
        private_users = True  # gaol default

    gpu = detect_gpu(files["binds"])

    return {
        "config": {
            "name": name,
            "runtime": "nspawn",
            "distro": distro,
            "image": None,
            "profiles": [],
            "network": {
                "mode": mode,
                "ports": ports,
                "dns": [],
                "hostname": None,
            },
            "resources": {
                "memory_mb": None,
                "cpu_shares": None,
                "cpu_count": None,
                "disk_mb": None,
            },
            "volumes": volumes,
            "environment": environment,
            "working_dir": None,
            "command": None,
            "privileged": False,
            "read_only": False,
            "private_users": private_users,
            "gpu": gpu,
            "devices": [],
            "users": [],
            "services": [],
        },
        "container_id": name,
        "runtime": "nspawn",
        "created_at": created_at.isoformat(),
        "updated_at": now.isoformat(),
        "last_started_at": None,
        "setup_done": True,
    }


def build_gaol_json(
    name: str,
    config: dict,
    distro: str,
    created_at: datetime,
) -> dict:
    """Build .gaol.json metadata matching nspawn runtime's _save_metadata format."""
    exec_cfg = config["exec"]
    files = config["files"]

    mode = determine_network_mode(config["network"])

    environment: dict[str, str] = {}
    for key, val in exec_cfg.get("_environments", []):
        environment[key] = val

    volumes: dict[str, str] = {}
    for host_path, container_path in files["binds"] + files["bind_ro"]:
        volumes[host_path] = container_path

    return {
        "name": name,
        "runtime": "nspawn",
        "distro": distro,
        "profiles": [],
        "created_at": created_at.isoformat(),
        "network_mode": mode,
        "environment": environment,
        "volumes": volumes,
        "privileged": False,
        "users": [],
        "services": [],
    }


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Adopt existing nspawn containers into the gaol store"
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print what would be written without writing",
    )
    parser.add_argument(
        "--skip",
        nargs="*",
        default=list(DEFAULT_SKIP),
        help=f"Container names to skip (default: {', '.join(sorted(DEFAULT_SKIP))})",
    )
    args = parser.parse_args()

    skip = set(args.skip)

    if os.geteuid() != 0:
        print(
            "Error: must be run as root (need access to /var/lib/machines/)",
            file=sys.stderr,
        )
        sys.exit(1)

    timestamps = get_image_timestamps()
    now = datetime.now()
    store_dir = get_store_dir()
    store_dir.mkdir(parents=True, exist_ok=True)

    nspawn_files = sorted(NSPAWN_DIR.glob("*.nspawn"))
    existing_store = {p.stem for p in store_dir.glob("*.yaml")}

    adopted: list[str] = []
    skipped: list[tuple[str, str]] = []
    errors: list[tuple[str, str]] = []

    for nspawn_file in nspawn_files:
        name = nspawn_file.stem

        if name in skip:
            skipped.append((name, "in skip list"))
            continue

        if name in existing_store:
            skipped.append((name, "already in store"))
            continue

        machine_dir = MACHINES_DIR / name
        if not machine_dir.exists():
            skipped.append((name, "no machine directory"))
            continue

        try:
            config = parse_nspawn_config(nspawn_file)
            distro = read_distro(name)
            created_at = timestamps.get(name, now)

            store_data = build_store_yaml(name, config, distro, created_at, now)
            gaol_json = build_gaol_json(name, config, distro, created_at)

            store_path = store_dir / f"{name}.yaml"
            json_path = machine_dir / ".gaol.json"

            mode = determine_network_mode(config["network"])
            caps = config["exec"].get("_capabilities", [])
            gpu = detect_gpu(config["files"]["binds"])

            if args.dry_run:
                print(f"\n  {name}")
                print(f"    distro={distro}  network={mode}  created={created_at.isoformat()}")
                if caps:
                    print(f"    capabilities: {' '.join(caps)}")
                if gpu:
                    print(f"    gpu: {gpu['vendor']}")
                print(f"    -> {store_path}")
                print(f"    -> {json_path}")
            else:
                store_path.write_text(
                    yaml.dump(store_data, default_flow_style=False, sort_keys=False)
                )
                json_path.write_text(json.dumps(gaol_json, indent=2, default=str))

            adopted.append(name)

        except Exception as e:
            errors.append((name, str(e)))

    # Summary
    label = "DRY RUN — " if args.dry_run else ""
    print(f"\n{label}Adopted {len(adopted)}, skipped {len(skipped)}, errors {len(errors)}")

    if skipped:
        print("\nSkipped:")
        for name, reason in skipped:
            print(f"  {name}: {reason}")

    if errors:
        print("\nErrors:")
        for name, err in errors:
            print(f"  {name}: {err}")


if __name__ == "__main__":
    main()
