"""VMOrchestrator: manage isolated dev environments via CloneBox or libvirt.

The class is composed from mixins, each grouping related methods:
  _InfraMixin      — _run, check_libvirt_available, _generate_domain_xml
  _LifecycleMixin  — create_vm, start_vm, stop_vm, pause_vm, resume_vm + backend helpers
  _SnapshotMixin   — snapshot_vm, restore_snapshot, delete_vm
  _StateMixin      — get_status, list_vms, _refresh_state_*, _get_vm_ip
  _ProjectMixin    — switch_project, _remount_project
"""

from __future__ import annotations

import json
import subprocess
import time
from dataclasses import asdict
from pathlib import Path
from typing import Any

import structlog

from proxym.virt._check import _check_clonebox
from proxym.virt._config import VMConfig, VMStatus
from proxym.virt._enums import VMBackend, VMState

logger = structlog.get_logger()


# ── Infra mixin ──────────────────────────────────────────────

class _InfraMixin:
    """Low-level helpers: shell commands and libvirt XML generation."""

    @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

    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


# ── Lifecycle mixin ──────────────────────────────────────────

class _LifecycleMixin:
    """VM creation, start, stop, pause, resume."""

    def create_vm(self, config: VMConfig) -> VMConfig:
        """Create a new VM with overlay disk and virtiofs mount.

        Uses CloneBox if available, otherwise falls back to virsh.
        Automatically mounts projects_root via virtiofs.
        """
        if not config.id:
            config.id = f"proxym-{config.name}-{int(time.time()) % 10000}"

        # Default code mount to projects_root
        if not config.code_mount_host:
            config.code_mount_host = str(self.projects_root)

        backend = self._resolve_backend(config)
        config.backend = backend.value

        if backend == VMBackend.CLONEBOX:
            self._create_vm_clonebox(config)
        else:
            self._create_vm_virsh(config)

        config.state = VMState.STOPPED
        self.vms[config.id] = config
        self._save_config(config)
        logger.info("vm_created", id=config.id, name=config.name, backend=backend.value)
        return config

    def _create_vm_clonebox(self, config: VMConfig) -> None:
        """Create VM via clonebox CLI."""
        cb = self._get_clonebox()
        if cb is None:
            raise RuntimeError("CloneBox not available")

        profile_path = cb.get_profile_path(config.tool) if config.tool else None
        extra: dict[str, Any] = {}
        if config.memory_mb:
            extra["memory"] = str(config.memory_mb)
        if config.vcpus:
            extra["vcpus"] = str(config.vcpus)

        result = cb.create_vm(config.id, profile_path=profile_path, extra_args=extra)
        if not result.success:
            logger.error("clonebox_create_failed", error=result.stderr)

    def _create_vm_virsh(self, config: VMConfig) -> None:
        """Create VM via direct virsh/qemu-img commands (fallback)."""
        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)

        self._generate_domain_xml(config)

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

        if vm.backend == VMBackend.CLONEBOX.value:
            return self._start_vm_clonebox(vm, vm_id)

        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 _start_vm_clonebox(self, vm: VMConfig, vm_id: str) -> bool:
        """Start VM via clonebox backend."""
        cb = self._get_clonebox()
        if not cb:
            return False
        result = cb.start_vm(vm.id)
        if result.success:
            vm.state = VMState.RUNNING
            vm.last_started_at = time.time()
            self._save_config(vm)
            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

        if vm.backend == VMBackend.CLONEBOX.value:
            return self._stop_vm_clonebox(vm)

        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 _stop_vm_clonebox(self, vm: VMConfig) -> bool:
        """Stop VM via clonebox backend."""
        cb = self._get_clonebox()
        if not cb:
            return False
        result = cb.stop_vm(vm.id)
        if result.success:
            vm.state = VMState.STOPPED
            self._save_config(vm)
            return True
        return False

    def pause_vm(self, vm_id: str) -> bool:
        """Pause (suspend) a VM."""
        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:
        """Resume a paused VM."""
        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


# ── Snapshot mixin ───────────────────────────────────────────

class _SnapshotMixin:
    """VM snapshots: create, restore, delete VM."""

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

        if vm.backend == VMBackend.CLONEBOX.value:
            return self._snapshot_clonebox(vm, snap_name)

        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 _snapshot_clonebox(self, vm: VMConfig, snap_name: str) -> bool:
        """Snapshot via clonebox backend."""
        cb = self._get_clonebox()
        if not cb:
            return False
        result = cb.snapshot(vm.id, snap_name)
        if result.success:
            vm.snapshot_name = snap_name
            self._save_config(vm)
            return True
        return False

    def restore_snapshot(self, vm_id: str, snapshot_name: str = "") -> bool:
        """Restore a previously saved snapshot."""
        vm = self.vms.get(vm_id)
        if not vm:
            return False
        snap = snapshot_name or vm.snapshot_name
        if not snap:
            return False

        if vm.backend == VMBackend.CLONEBOX.value:
            cb = self._get_clonebox()
            if cb:
                return cb.restore_snapshot(vm.id, snap).success

        result = self._run(f"virsh snapshot-revert {vm.id} {snap}", check=False)
        return result.returncode == 0

    def delete_vm(self, vm_id: str) -> bool:
        """Delete a VM and all its resources."""
        vm = self.vms.get(vm_id)
        if not vm:
            return False

        if vm.backend == VMBackend.CLONEBOX.value:
            cb = self._get_clonebox()
            if cb:
                cb.delete_vm(vm.id)
        else:
            self.stop_vm(vm_id, force=True)
            self._run(f"virsh undefine {vm.id} --snapshots-metadata", check=False)
            overlay = Path(vm.overlay_path)
            if overlay.exists():
                overlay.unlink()

        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


# ── State mixin ──────────────────────────────────────────────

class _StateMixin:
    """VM status queries and state refresh."""

    def get_status(self, vm_id: str) -> VMStatus | None:
        """Get current status of a VM, querying the backend for live state."""
        vm = self.vms.get(vm_id)
        if not vm:
            return None

        if vm.backend == VMBackend.CLONEBOX.value:
            self._refresh_state_clonebox(vm)
        else:
            self._refresh_state_virsh(vm)

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

        ip = self._get_vm_ip(vm.id) if vm.state == VMState.RUNNING else ""

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

    def _refresh_state_clonebox(self, vm: VMConfig) -> None:
        """Update VM state from clonebox health output."""
        cb = self._get_clonebox()
        if not cb:
            return
        health = cb.health(vm.id)
        if not health.success:
            return
        out = health.stdout.lower()
        if "running" in out:
            vm.state = VMState.RUNNING
        elif "stopped" in out or "shut off" in out:
            vm.state = VMState.STOPPED
        elif "paused" in out:
            vm.state = VMState.PAUSED

    def _refresh_state_virsh(self, vm: VMConfig) -> None:
        """Update VM state from virsh domstate."""
        result = self._run(f"virsh domstate {vm.id}", check=False)
        if result.returncode != 0:
            return
        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)

    def _get_vm_ip(self, vm_id: str) -> str:
        """Extract VM IP address from virsh DHCP leases."""
        ip_result = self._run(
            f"virsh domifaddr {vm_id} --source agent", check=False
        )
        if ip_result.returncode != 0:
            return ""
        for line in ip_result.stdout.splitlines():
            if "ipv4" in line:
                for p in line.split():
                    if "/" in p and "." in p:
                        return p.split("/")[0]
        return ""

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


# ── Project mixin ────────────────────────────────────────────

class _ProjectMixin:
    """Project switching (hot-swap code mounts)."""

    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:
            self._remount_project(vm, vm_id)

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

    def _remount_project(self, vm: VMConfig, vm_id: str) -> None:
        """SSH into VM and remount shared filesystem."""
        status = self.get_status(vm_id)
        if not status or not status.ip_address:
            return
        self._run(
            f"ssh -o StrictHostKeyChecking=no dev@{status.ip_address} "
            f"'sudo mount -t 9p -o trans=virtio code_share {vm.code_mount_guest}'",
            check=False,
        )


# ── VMOrchestrator (composition) ─────────────────────────────

class VMOrchestrator(
    _InfraMixin, _LifecycleMixin, _SnapshotMixin, _StateMixin, _ProjectMixin,
):
    """Manages VMs for isolated dev environments.

    Backend selection:
      - CloneBox (preferred): uses clonebox CLI with YAML profiles
      - virsh (fallback): direct libvirt virsh commands

    Set backend="auto" (default) to auto-detect.
    """

    def __init__(self, config_dir: str = "~/.config/proxym/vms",
                 images_dir: str = "~/.local/share/proxym/images",
                 overlays_dir: str = "~/.local/share/proxym/overlays",
                 profiles_dir: str = "~/.config/proxym/profiles",
                 projects_root: str = "~/github"):
        self.config_dir = Path(config_dir).expanduser()
        self.images_dir = Path(images_dir).expanduser()
        self.overlays_dir = Path(overlays_dir).expanduser()
        self.profiles_dir = Path(profiles_dir).expanduser()
        self.projects_root = Path(projects_root).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._clonebox = None
        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))

    def _get_clonebox(self):
        """Get or create the CloneBox adapter (lazy init)."""
        if self._clonebox is None and _check_clonebox():
            from proxym.virt.clonebox_adapter import CloneBoxAdapter
            self._clonebox = CloneBoxAdapter(
                profiles_dir=str(self.profiles_dir),
            )
        return self._clonebox

    def _resolve_backend(self, config: VMConfig) -> VMBackend:
        """Determine which backend to use for a VM."""
        if config.backend == "clonebox":
            return VMBackend.CLONEBOX
        if config.backend == "virsh":
            return VMBackend.VIRSH
        # auto: prefer clonebox if available
        if self._get_clonebox() is not None:
            return VMBackend.CLONEBOX
        return VMBackend.VIRSH
