"""GPU hardware detection with graceful CPU fallback."""

import json
import os
import shutil
import subprocess
import time


def _detect_ram_mb() -> int:
    """Return host physical RAM in MB, or 0 if unavailable."""
    try:
        pages = os.sysconf("SC_PHYS_PAGES")
        page_size = os.sysconf("SC_PAGE_SIZE")
        return int((pages * page_size) // (1024 * 1024))
    except (AttributeError, OSError, ValueError):
        pass
    return _detect_windows_ram_mb()


def _detect_windows_ram_mb() -> int:
    """Return Windows physical RAM in MB, or 0."""
    try:
        import ctypes

        class _MemoryStatusEx(ctypes.Structure):
            _fields_ = [
                ("dwLength", ctypes.c_ulong),
                ("dwMemoryLoad", ctypes.c_ulong),
                ("ullTotalPhys", ctypes.c_ulonglong),
                ("ullAvailPhys", ctypes.c_ulonglong),
                ("ullTotalPageFile", ctypes.c_ulonglong),
                ("ullAvailPageFile", ctypes.c_ulonglong),
                ("ullTotalVirtual", ctypes.c_ulonglong),
                ("ullAvailVirtual", ctypes.c_ulonglong),
                ("ullAvailExtendedVirtual", ctypes.c_ulonglong),
            ]

        status = _MemoryStatusEx()
        status.dwLength = ctypes.sizeof(_MemoryStatusEx)
        if ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(status)):
            return int(status.ullTotalPhys // (1024 * 1024))
    except (AttributeError, OSError, ValueError):
        pass

    try:
        result = subprocess.run(
            [
                "powershell",
                "-NoProfile",
                "-Command",
                "(Get-CimInstance Win32_ComputerSystem).TotalPhysicalMemory",
            ],
            capture_output=True,
            text=True,
            timeout=5,
        )
        if result.returncode == 0 and result.stdout.strip():
            return int(result.stdout.strip()) // (1024 * 1024)
    except (FileNotFoundError, subprocess.TimeoutExpired, ValueError):
        pass
    return 0


def _detect_windows_gpu_memory() -> dict | None:
    """Return Windows GPU memory metadata from Win32_VideoController, if available."""
    try:
        result = subprocess.run(
            [
                "powershell",
                "-NoProfile",
                "-Command",
                (
                    "Get-CimInstance Win32_VideoController | "
                    "Select-Object Name,AdapterRAM | ConvertTo-Json -Compress"
                ),
            ],
            capture_output=True,
            text=True,
            timeout=5,
        )
    except (FileNotFoundError, subprocess.TimeoutExpired):
        return None

    if result.returncode != 0 or not result.stdout.strip():
        return None
    try:
        raw = json.loads(result.stdout)
    except json.JSONDecodeError:
        return None
    entries = raw if isinstance(raw, list) else [raw]
    best: dict | None = None
    for entry in entries:
        if not isinstance(entry, dict):
            continue
        name = str(entry.get("Name") or "").strip()
        if not name:
            continue
        try:
            adapter_ram = int(entry.get("AdapterRAM") or 0)
        except (TypeError, ValueError):
            adapter_ram = 0
        vram_mb = max(0, adapter_ram // (1024 * 1024))
        if best is None or vram_mb > best["vram_mb"]:
            best = {"gpu_name": name, "vram_mb": vram_mb}
    return best


def _detect_nvidia_smi_gpu_memory() -> dict | None:
    """Return NVIDIA GPU memory metadata from nvidia-smi, if available."""
    try:
        result = subprocess.run(
            ["nvidia-smi", "--query-gpu=name,memory.total", "--format=csv,noheader,nounits"],
            capture_output=True, text=True, timeout=5,
        )
        if result.returncode == 0 and result.stdout.strip():
            line = result.stdout.strip().splitlines()[0]
            parts = line.split(",", 1)
            gpu_name = parts[0].strip()
            vram_mb = int(parts[1].strip()) if len(parts) > 1 else 0
            return {"gpu_name": gpu_name, "vram_mb": max(0, vram_mb)}
    except (FileNotFoundError, subprocess.TimeoutExpired, ValueError, IndexError):
        pass
    return None


def _detect_torch_cuda_inventory(torch_module) -> dict | None:
    """Return torch-visible CUDA/HIP GPU metadata without running kernels."""
    try:
        if not torch_module.cuda.is_available() or torch_module.cuda.device_count() < 1:
            return None
        idx = torch_module.cuda.current_device()
        name = torch_module.cuda.get_device_name(idx)
        props = torch_module.cuda.get_device_properties(idx)
        vram_mb = int(props.total_memory // (1024 * 1024))
        gpu = {"gpu_name": name, "vram_mb": max(0, vram_mb)}
        gcn_arch = getattr(props, "gcnArchName", None)
        if gcn_arch:
            gpu["gcn_arch"] = str(gcn_arch)
        return gpu
    except Exception:
        return None


def _torch_cuda_is_executable(torch_module) -> bool:
    """Return True only if this Python process can execute a CUDA tensor op."""
    try:
        if not torch_module.cuda.is_available():
            return False
        probe = torch_module.empty((1,), device="cuda")
        probe += 1
        torch_module.cuda.synchronize()
        return True
    except Exception:
        return False


def _gpu_inventory_profile(gpu: dict | None, ram_mb: int) -> dict | None:
    if gpu is None:
        return None
    profile = {
        "device": "cpu",
        "gpu_name": gpu["gpu_name"],
        "vram_mb": gpu["vram_mb"],
        "dedicated_vram_mb": gpu["vram_mb"],
        "shared_vram_mb": max(0, ram_mb // 2),
        "ram_mb": ram_mb,
        "cuda_available": False,
    }
    if gpu.get("gcn_arch"):
        profile["gcn_arch"] = gpu["gcn_arch"]
    return profile


def with_forced_cpu(hw: dict) -> dict:
    """Return a hardware profile forced to CPU execution.

    Keeps detected GPU metadata for diagnostics and tracker registration context,
    but clears CUDA availability so startup and the model backend stay on CPU.
    """
    forced = dict(hw)
    forced["device"] = "cpu"
    forced["cuda_available"] = False
    return forced


def _with_model_drive(profile: dict) -> dict:
    """Attach free space for the default model cache drive to tracker diagnostics."""
    try:
        cache_root = os.path.expanduser("~/.cache/meshnet/shards")
        profile["model_drive_free_bytes"] = shutil.disk_usage(os.path.expanduser("~")).free
        profile["model_drive_path"] = cache_root
    except OSError:
        pass
    return profile


def detect_hardware() -> dict:
    """Detect GPU model and available VRAM. Returns hardware profile dict."""
    ram_mb = _detect_ram_mb()
    torch_gpu: dict | None = None
    try:
        import torch  # type: ignore[import]

        torch_gpu = _detect_torch_cuda_inventory(torch)
        if _torch_cuda_is_executable(torch):
            if torch_gpu is None:
                torch_gpu = _detect_torch_cuda_inventory(torch)
            name = torch_gpu["gpu_name"] if torch_gpu is not None else "CUDA GPU"
            vram_mb = torch_gpu["vram_mb"] if torch_gpu is not None else 0
            shared_vram_mb = max(0, ram_mb // 2)
            profile = {
                "device": "cuda",
                "gpu_name": name,
                "vram_mb": vram_mb,
                "dedicated_vram_mb": vram_mb,
                "shared_vram_mb": shared_vram_mb,
                "ram_mb": ram_mb,
                "cuda_available": True,
            }
            if torch_gpu is not None and torch_gpu.get("gcn_arch"):
                profile["gcn_arch"] = torch_gpu["gcn_arch"]
            return _with_model_drive(profile)
    except ImportError:
        pass

    torch_inventory = _gpu_inventory_profile(torch_gpu, ram_mb)
    if torch_inventory is not None:
        return _with_model_drive(torch_inventory)

    nvidia_gpu = _gpu_inventory_profile(_detect_nvidia_smi_gpu_memory(), ram_mb)
    if nvidia_gpu is not None:
        return _with_model_drive(nvidia_gpu)

    windows_gpu = _gpu_inventory_profile(_detect_windows_gpu_memory(), ram_mb)
    if windows_gpu is not None:
        return _with_model_drive(windows_gpu)

    return _with_model_drive({
        "device": "cpu",
        "gpu_name": None,
        "vram_mb": 0,
        "dedicated_vram_mb": 0,
        "shared_vram_mb": 0,
        "ram_mb": ram_mb,
        "cuda_available": False,
    })


def benchmark_throughput_checked(device_str: str = "cpu") -> tuple[float, bool, str | None]:
    """
    Estimate compute throughput via a synthetic transformer GEMM benchmark.

    Runs hidden_size × (hidden_size*4) matmul — the dominant op in FFN layers —
    and returns iterations/second as a relative speed index.  Higher = faster.
    The value is used as benchmark_tokens_per_sec in tracker registration for
    routing tiebreaks; it is not an absolute token rate.

    Returns (score, ok, error). Score falls back to 1.0 when the requested
    device cannot run the benchmark.
    """
    try:
        import torch  # type: ignore[import]

        device = torch.device(device_str)
        # bfloat16 on CUDA matches real inference dtype; float32 on CPU avoids
        # precision-downcast surprises on older hardware without bfloat16 support.
        dtype = torch.bfloat16 if device_str == "cuda" else torch.float32

        # hidden_size=2048 is representative of a mid-sized model; large enough
        # that BLAS finds an efficient kernel on both GPU and CPU.
        hidden_size = 2048
        a = torch.randn(1, hidden_size, dtype=dtype, device=device)
        b = torch.randn(hidden_size, hidden_size * 4, dtype=dtype, device=device)

        def _sync() -> None:
            if device_str == "cuda":
                torch.cuda.synchronize()

        # Warmup: prime caches and JIT compilation.
        for _ in range(10):
            torch.matmul(a, b)
        _sync()

        n_iters = 50
        t0 = time.perf_counter()
        for _ in range(n_iters):
            torch.matmul(a, b)
        _sync()
        elapsed = time.perf_counter() - t0

        return round(n_iters / max(elapsed, 1e-9), 2), True, None
    except Exception as exc:
        return 1.0, False, f"{type(exc).__name__}: {exc}"


def benchmark_throughput(device_str: str = "cpu") -> float:
    """Return only the numeric throughput index, preserving the legacy API."""
    score, _ok, _error = benchmark_throughput_checked(device_str)
    return score