"""KVM kernel module management tools (kvm_unload, kvm_reload).

These tools manage KVM kernel modules to resolve the KVM/VirtualBox conflict.
KVM modules must be unloaded before VirtualBox can use hardware virtualization.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Annotated, Literal

from fastmcp.tools import tool
from pydantic import Field

from mcp_vm_blackbox._helpers import _resolve, _ssh_run
from mcp_vm_blackbox.models import KvmReloadResult, KvmUnloadResult

if TYPE_CHECKING:
    from mcp_vm_blackbox.targets import TargetConfig


async def _lsmod_kvm(cfg: TargetConfig) -> list[str]:
    """List loaded KVM modules via lsmod.

    Args:
        cfg: Resolved TargetConfig.

    Returns:
        List of KVM module names (e.g. ['kvm_intel', 'kvm'] or ['kvm_amd', 'kvm']).
    """
    result = await _ssh_run(cfg, ["bash", "-c", "lsmod | grep '^kvm' | awk '{print $1}'"])
    if result.returncode != 0:
        return []
    lines = result.stdout.strip().splitlines()
    return [line.strip() for line in lines if line.strip()]


async def _detect_cpu_vendor(cfg: TargetConfig) -> Literal["intel", "amd", "unknown"]:
    """Detect CPU vendor via /proc/cpuinfo or lscpu.

    Args:
        cfg: Resolved TargetConfig.

    Returns:
        'intel', 'amd', or 'unknown'.
    """
    # Try lscpu first (more reliable)
    result = await _ssh_run(
        cfg,
        ["bash", "-c", "lscpu 2>/dev/null | grep -i 'vendor id' || cat /proc/cpuinfo | grep -i 'vendor_id' | head -1"],
    )
    if result.returncode == 0 and result.stdout.strip():
        vendor_line = result.stdout.strip().lower()
        if "genuineintel" in vendor_line or "intel" in vendor_line:
            return "intel"
        if "authenticamd" in vendor_line or "amd" in vendor_line:
            return "amd"

    # Fallback: check for specific CPU flags
    result = await _ssh_run(cfg, ["bash", "-c", "cat /proc/cpuinfo | grep -E 'flags|model name' | head -4"])
    if result.returncode == 0:
        cpuinfo = result.stdout.lower()
        if "intel" in cpuinfo:
            return "intel"
        if "amd" in cpuinfo:
            return "amd"

    return "unknown"


async def _check_kvm_users(cfg: TargetConfig) -> list[int]:
    """Check for processes using /dev/kvm via fuser.

    Args:
        cfg: Resolved TargetConfig.

    Returns:
        List of PIDs using /dev/kvm, empty if none or /dev/kvm absent.
    """
    result = await _ssh_run(cfg, ["bash", "-c", "fuser /dev/kvm 2>/dev/null || true"])
    if result.returncode != 0 or not result.stdout.strip():
        return []
    # fuser output is space-separated PIDs
    try:
        pids = [int(pid) for pid in result.stdout.split() if pid.strip()]
    except ValueError:
        return []
    return pids


def _get_vendor_module(cpu_vendor: str, modules: list[str]) -> str | None:
    """Determine the vendor module name based on CPU vendor and loaded modules.

    Args:
        cpu_vendor: Detected CPU vendor ('intel', 'amd', or 'unknown').
        modules: List of currently loaded KVM modules.

    Returns:
        Vendor module name or None if unknown.
    """
    if cpu_vendor == "intel":
        return "kvm_intel"
    if cpu_vendor == "amd":
        return "kvm_amd"
    # Unknown vendor - try to infer from loaded modules
    if "kvm_intel" in modules:
        return "kvm_intel"
    if "kvm_amd" in modules:
        return "kvm_amd"
    return None


async def _unload_modules(cfg: TargetConfig, modules: list[str], vendor_module: str | None) -> tuple[bool, str]:
    """Unload KVM modules in correct order.

    Args:
        cfg: Resolved TargetConfig.
        modules: List of currently loaded KVM modules.
        vendor_module: Vendor-specific module name.

    Returns:
        Tuple of (success, error_message).
    """
    unload_order = []
    if vendor_module and vendor_module in modules:
        unload_order.append(vendor_module)
    if "kvm" in modules:
        unload_order.append("kvm")

    for module in unload_order:
        result = await _ssh_run(cfg, ["modprobe", "-r", module])
        if result.returncode != 0:
            # Check if already unloaded
            current = await _lsmod_kvm(cfg)
            if module in current:
                return (False, f"Failed to unload {module}: {result.stderr.strip() or 'unknown error'}")
    return True, ""


@tool
async def kvm_unload(
    target: Annotated[str, Field(description='Named target or raw SSH string. Defaults to "local".')] = "local",
    dry_run: Annotated[bool, Field(description="If true, report planned actions without executing.")] = False,
) -> KvmUnloadResult:
    """Unload KVM kernel modules to resolve VirtualBox KVM conflict.

    Detects CPU vendor, checks for active KVM users, and unloads modules
    in correct dependency order (vendor module before kvm core).

    Args:
        target: Named target or raw SSH string.
        dry_run: If True, return planned actions without executing.

    Returns:
        KvmUnloadResult with status, modules_before, modules_after, cpu_vendor,
        and error if blocked or failed.
    """
    cfg = _resolve(target)

    # Detect CPU vendor and current state
    cpu_vendor = await _detect_cpu_vendor(cfg)
    modules_before = await _lsmod_kvm(cfg)

    if not modules_before:
        # Already unloaded
        return KvmUnloadResult(status="ok", modules_before=[], modules_after=[], cpu_vendor=cpu_vendor, dry_run=dry_run)

    # Check for active KVM users
    active_pids = await _check_kvm_users(cfg)
    if active_pids:
        return KvmUnloadResult(
            status="blocked",
            modules_before=modules_before,
            modules_after=modules_before,
            cpu_vendor=cpu_vendor,
            dry_run=dry_run,
            error=f"Active KVM users detected (PIDs: {active_pids}). Stop VMs/QEMU before unloading.",
        )

    # Determine vendor module
    vendor_module = _get_vendor_module(cpu_vendor, modules_before)

    if dry_run:
        return KvmUnloadResult(
            status="ok", modules_before=modules_before, modules_after=[], cpu_vendor=cpu_vendor, dry_run=True
        )

    # Unload modules
    success, error = await _unload_modules(cfg, modules_before, vendor_module)
    if not success:
        return KvmUnloadResult(
            status="error",
            modules_before=modules_before,
            modules_after=await _lsmod_kvm(cfg),
            cpu_vendor=cpu_vendor,
            dry_run=False,
            error=error,
        )

    # Verify final state
    modules_after = await _lsmod_kvm(cfg)
    return KvmUnloadResult(
        status="ok", modules_before=modules_before, modules_after=modules_after, cpu_vendor=cpu_vendor, dry_run=False
    )


@tool
async def kvm_reload(
    target: Annotated[str, Field(description='Named target or raw SSH string. Defaults to "local".')] = "local",
) -> KvmReloadResult:
    """Reload KVM kernel modules after VirtualBox use.

    Loads KVM core module and the appropriate vendor module (kvm_intel or kvm_amd)
    based on detected CPU vendor.

    Args:
        target: Named target or raw SSH string.

    Returns:
        KvmReloadResult with status, modules_before, modules_after, cpu_vendor,
        and error if failed.
    """
    cfg = _resolve(target)

    # Detect CPU vendor
    cpu_vendor = await _detect_cpu_vendor(cfg)

    # Get current KVM modules
    modules_before = await _lsmod_kvm(cfg)

    # Determine vendor module name
    if cpu_vendor == "intel":
        vendor_module = "kvm_intel"
    elif cpu_vendor == "amd":
        vendor_module = "kvm_amd"
    else:
        return KvmReloadResult(
            status="error",
            modules_before=modules_before,
            modules_after=modules_before,
            cpu_vendor=cpu_vendor,
            error="Cannot reload KVM: unknown CPU vendor (expected 'intel' or 'amd')",
        )

    # Check if already loaded
    if "kvm" in modules_before and vendor_module in modules_before:
        return KvmReloadResult(
            status="ok", modules_before=modules_before, modules_after=modules_before, cpu_vendor=cpu_vendor
        )

    # Load kvm core first, then vendor module
    load_order = ["kvm", vendor_module]

    for module in load_order:
        if module in modules_before:
            # Already loaded, skip
            continue
        result = await _ssh_run(cfg, ["modprobe", module])
        if result.returncode != 0:
            return KvmReloadResult(
                status="error",
                modules_before=modules_before,
                modules_after=await _lsmod_kvm(cfg),
                cpu_vendor=cpu_vendor,
                error=f"Failed to load {module}: {result.stderr.strip() or 'unknown error'}",
            )

    # Verify final state
    modules_after = await _lsmod_kvm(cfg)
    return KvmReloadResult(
        status="ok", modules_before=modules_before, modules_after=modules_after, cpu_vendor=cpu_vendor
    )
