"""VM Orchestrator: manage libvirt-based dev environments.

Each tool/account combo gets an isolated VM with:
  - Dedicated browser profile (preserved cache/sessions)
  - Shared code directory (9p/virtiofs mount from host)
  - Network access through the AI proxy
  - Snapshot/restore for quick switching between projects

Architecture:
  Host (code + proxym)
    └── libvirt VMs
        ├── vm-windsurf-work   (Windsurf + Chrome profile A)
        ├── vm-windsurf-oss    (Windsurf + Chrome profile B)
        ├── vm-cursor-client1  (Cursor + separate account)
        └── vm-vscode-dev      (VS Code + Roo Code)

Code is on the host, mounted read-write via virtiofs.
Browser profiles are inside the VM (isolated per account).
"""

from __future__ import annotations

import json
import subprocess
import time
from dataclasses import asdict, dataclass, field
from enum import Enum
from pathlib import Path

import structlog

logger = structlog.get_logger()


class VMState(str, Enum):
    STOPPED = "stopped"
    RUNNING = "running"
    PAUSED = "paused"
    ERROR = "error"
    CREATING = "creating"


@dataclass
class VMConfig:
    """Configuration for a development VM."""
    id: str = ""
    name: str = ""                       # e.g., "windsurf-work"
    account_id: str = ""                 # linked account
    tool: str = ""                       # tool type

    # Resources
    vcpus: int = 2
    memory_mb: int = 4096
    disk_gb: int = 20

    # Base image
    base_image: str = "fedora-40-dev"    # qcow2 base
    overlay_path: str = ""               # cow overlay per VM

    # Mounts: host → VM
    code_mount_host: str = ""            # e.g., /home/user/projects
    code_mount_guest: str = "/code"      # mount point in VM

    # Networking
    proxy_address: str = "10.0.0.1:4000"  # AI proxy on host
    forward_ports: list[int] = field(default_factory=lambda: [3000, 8080])

    # Browser profile
    browser_profile_path: str = ""        # inside VM, persisted as overlay
    preserve_browser_cache: bool = True

    # State
    state: VMState = VMState.STOPPED
    created_at: float = field(default_factory=time.time)
    last_started_at: float = 0.0
    snapshot_name: str = ""


@dataclass
class VMStatus:
    """Runtime status of a VM."""
    id: str
    name: str
    state: VMState
    account_id: str
    tool: str
    uptime_seconds: float = 0
    cpu_usage_pct: float = 0
    memory_used_mb: int = 0
    code_mount: str = ""
    ip_address: str = ""


class VMOrchestrator:
    """Manages libvirt VMs for isolated dev environments.

    Uses virsh CLI as the interface to libvirt — avoids the fragile
    libvirt Python bindings while keeping full functionality.
    """

    def __init__(self, config_dir: str = "~/.config/proxym/vms",
                 images_dir: str = "~/.local/share/proxym/images",
                 overlays_dir: str = "~/.local/share/proxym/overlays"):
        self.config_dir = Path(config_dir).expanduser()
        self.images_dir = Path(images_dir).expanduser()
        self.overlays_dir = Path(overlays_dir).expanduser()
        self.config_dir.mkdir(parents=True, exist_ok=True)
        self.overlays_dir.mkdir(parents=True, exist_ok=True)
        self.vms: dict[str, VMConfig] = {}
        self._load_configs()

    def _load_configs(self) -> None:
        for f in self.config_dir.glob("*.json"):
            try:
                data = json.loads(f.read_text())
                vm = VMConfig(**{k: v for k, v in data.items()
                                if k in VMConfig.__dataclass_fields__})
                vm.state = VMState(data.get("state", "stopped"))
                self.vms[vm.id] = vm
            except Exception as e:
                logger.warning("vm_config_load_failed", file=f.name, error=str(e))

    def _save_config(self, vm: VMConfig) -> None:
        path = self.config_dir / f"{vm.id}.json"
        path.write_text(json.dumps(asdict(vm), indent=2, default=str))

    # ── VM Lifecycle ─────────────────────────────────────────

    def create_vm(self, config: VMConfig) -> VMConfig:
        """Create a new VM with overlay disk and virtiofs mount."""
        if not config.id:
            config.id = f"aiproxy-{config.name}-{int(time.time()) % 10000}"

        # Create overlay disk
        overlay = self.overlays_dir / f"{config.id}.qcow2"
        base = self.images_dir / f"{config.base_image}.qcow2"

        if base.exists():
            self._run(f"qemu-img create -f qcow2 -b {base} -F qcow2 {overlay}")
            config.overlay_path = str(overlay)
        else:
            logger.warning("base_image_missing", image=config.base_image,
                          hint=f"Place {config.base_image}.qcow2 in {self.images_dir}")
            config.overlay_path = str(overlay)

        config.state = VMState.STOPPED
        self.vms[config.id] = config
        self._save_config(config)

        # Generate libvirt XML
        self._generate_domain_xml(config)

        logger.info("vm_created", id=config.id, name=config.name)
        return config

    def start_vm(self, vm_id: str) -> bool:
        """Start a VM."""
        vm = self.vms.get(vm_id)
        if not vm:
            return False

        result = self._run(f"virsh start {vm.id}", check=False)
        if result.returncode == 0:
            vm.state = VMState.RUNNING
            vm.last_started_at = time.time()
            self._save_config(vm)
            logger.info("vm_started", id=vm_id)
            return True

        logger.error("vm_start_failed", id=vm_id, error=result.stderr)
        return False

    def stop_vm(self, vm_id: str, force: bool = False) -> bool:
        """Stop a VM (graceful shutdown or force destroy)."""
        vm = self.vms.get(vm_id)
        if not vm:
            return False

        cmd = f"virsh destroy {vm.id}" if force else f"virsh shutdown {vm.id}"
        result = self._run(cmd, check=False)
        if result.returncode == 0:
            vm.state = VMState.STOPPED
            self._save_config(vm)
            logger.info("vm_stopped", id=vm_id, force=force)
            return True
        return False

    def pause_vm(self, vm_id: str) -> bool:
        vm = self.vms.get(vm_id)
        if not vm:
            return False
        result = self._run(f"virsh suspend {vm.id}", check=False)
        if result.returncode == 0:
            vm.state = VMState.PAUSED
            self._save_config(vm)
            return True
        return False

    def resume_vm(self, vm_id: str) -> bool:
        vm = self.vms.get(vm_id)
        if not vm:
            return False
        result = self._run(f"virsh resume {vm.id}", check=False)
        if result.returncode == 0:
            vm.state = VMState.RUNNING
            self._save_config(vm)
            return True
        return False

    def snapshot_vm(self, vm_id: str, name: str = "") -> bool:
        """Create a snapshot for quick restore."""
        vm = self.vms.get(vm_id)
        if not vm:
            return False
        snap_name = name or f"snap-{int(time.time())}"
        result = self._run(
            f"virsh snapshot-create-as {vm.id} {snap_name} --atomic",
            check=False,
        )
        if result.returncode == 0:
            vm.snapshot_name = snap_name
            self._save_config(vm)
            return True
        return False

    def restore_snapshot(self, vm_id: str, snapshot_name: str = "") -> bool:
        vm = self.vms.get(vm_id)
        if not vm:
            return False
        snap = snapshot_name or vm.snapshot_name
        if not snap:
            return False
        result = self._run(f"virsh snapshot-revert {vm.id} {snap}", check=False)
        return result.returncode == 0

    def delete_vm(self, vm_id: str) -> bool:
        vm = self.vms.get(vm_id)
        if not vm:
            return False
        self.stop_vm(vm_id, force=True)
        self._run(f"virsh undefine {vm.id} --snapshots-metadata", check=False)
        # Remove overlay
        overlay = Path(vm.overlay_path)
        if overlay.exists():
            overlay.unlink()
        # Remove config
        cfg = self.config_dir / f"{vm.id}.json"
        if cfg.exists():
            cfg.unlink()
        del self.vms[vm_id]
        logger.info("vm_deleted", id=vm_id)
        return True

    # ── Status & Listing ─────────────────────────────────────

    def get_status(self, vm_id: str) -> VMStatus | None:
        vm = self.vms.get(vm_id)
        if not vm:
            return None

        # Query actual state from libvirt
        result = self._run(f"virsh domstate {vm.id}", check=False)
        if result.returncode == 0:
            state_str = result.stdout.strip()
            state_map = {
                "running": VMState.RUNNING,
                "shut off": VMState.STOPPED,
                "paused": VMState.PAUSED,
            }
            vm.state = state_map.get(state_str, VMState.ERROR)

        uptime = 0.0
        if vm.state == VMState.RUNNING and vm.last_started_at:
            uptime = time.time() - vm.last_started_at

        # Get IP via DHCP lease
        ip = ""
        if vm.state == VMState.RUNNING:
            ip_result = self._run(
                f"virsh domifaddr {vm.id} --source agent", check=False
            )
            if ip_result.returncode == 0:
                for line in ip_result.stdout.splitlines():
                    if "ipv4" in line:
                        parts = line.split()
                        for p in parts:
                            if "/" in p and "." in p:
                                ip = p.split("/")[0]
                                break

        return VMStatus(
            id=vm.id,
            name=vm.name,
            state=vm.state,
            account_id=vm.account_id,
            tool=vm.tool,
            uptime_seconds=uptime,
            code_mount=vm.code_mount_host,
            ip_address=ip,
        )

    def list_vms(self) -> list[VMStatus]:
        return [s for vm_id in self.vms if (s := self.get_status(vm_id))]

    # ── Project Switching ────────────────────────────────────

    def switch_project(self, vm_id: str, new_project_path: str) -> bool:
        """Hot-swap the code mount to a different project directory.

        Uses virtiofs or 9p — requires VM to use shared filesystem.
        For simplicity, we update the config and remount.
        """
        vm = self.vms.get(vm_id)
        if not vm:
            return False

        old_path = vm.code_mount_host
        vm.code_mount_host = new_project_path
        self._save_config(vm)

        if vm.state == VMState.RUNNING:
            # For virtiofs, we'd need to update the shared memory backend
            # Simpler approach: SSH into VM and remount
            ip = ""
            status = self.get_status(vm_id)
            if status:
                ip = status.ip_address
            if ip:
                self._run(
                    f"ssh -o StrictHostKeyChecking=no dev@{ip} "
                    f"'sudo mount -t 9p -o trans=virtio code_share {vm.code_mount_guest}'",
                    check=False,
                )

        logger.info("project_switched", vm=vm_id, old=old_path, new=new_project_path)
        return True

    # ── Libvirt XML Generation ───────────────────────────────

    def _generate_domain_xml(self, vm: VMConfig) -> str:
        """Generate libvirt domain XML for the VM."""
        xml = f"""<domain type='kvm'>
  <name>{vm.id}</name>
  <memory unit='MiB'>{vm.memory_mb}</memory>
  <vcpu>{vm.vcpus}</vcpu>

  <os>
    <type arch='x86_64' machine='q35'>hvm</type>
    <boot dev='hd'/>
  </os>

  <features>
    <acpi/><apic/>
  </features>

  <cpu mode='host-passthrough'/>

  <devices>
    <!-- Overlay disk -->
    <disk type='file' device='disk'>
      <driver name='qemu' type='qcow2'/>
      <source file='{vm.overlay_path}'/>
      <target dev='vda' bus='virtio'/>
    </disk>

    <!-- Network (NAT with port forwards) -->
    <interface type='network'>
      <source network='default'/>
      <model type='virtio'/>
    </interface>

    <!-- Shared filesystem: host code → VM /code -->
    <filesystem type='mount' accessmode='passthrough'>
      <driver type='virtiofs'/>
      <source dir='{vm.code_mount_host}'/>
      <target dir='code_share'/>
    </filesystem>

    <!-- Console -->
    <serial type='pty'><target port='0'/></serial>
    <console type='pty'><target type='serial' port='0'/></console>

    <!-- VGA for GUI tools -->
    <graphics type='spice' autoport='yes'>
      <listen type='address' address='127.0.0.1'/>
    </graphics>
    <video><model type='virtio' heads='1' primary='yes'/></video>

    <!-- Shared memory for virtiofs -->
    <memballoon model='virtio'/>
  </devices>

  <memoryBacking>
    <source type='memfd'/>
    <access mode='shared'/>
  </memoryBacking>
</domain>"""

        xml_path = self.config_dir / f"{vm.id}.xml"
        xml_path.write_text(xml)

        # Define in libvirt
        self._run(f"virsh define {xml_path}", check=False)
        return xml

    # ── Helpers ───────────────────────────────────────────────

    @staticmethod
    def _run(cmd: str, check: bool = True) -> subprocess.CompletedProcess:
        try:
            return subprocess.run(
                cmd, shell=True, capture_output=True, text=True,
                timeout=30,
            )
        except subprocess.TimeoutExpired:
            return subprocess.CompletedProcess(cmd, 1, "", "timeout")
        except Exception as e:
            return subprocess.CompletedProcess(cmd, 1, "", str(e))

    @staticmethod
    def check_libvirt_available() -> dict:
        """Check if libvirt is available on this system."""
        checks = {}
        for cmd, label in [
            ("virsh version", "libvirt"),
            ("qemu-img --version", "qemu"),
            ("virt-install --version", "virt-install"),
        ]:
            try:
                result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=5)
                checks[label] = result.returncode == 0
            except Exception:
                checks[label] = False
        return checks
