"""
GPU data collection — NVML for device stats, Windows perf counters for
per-process VRAM, psutil for process name resolution.
"""

from __future__ import annotations

import os
import re
import subprocess
from dataclasses import dataclass, asdict
from typing import Optional

try:
    import psutil as _psutil
except ImportError:
    _psutil = None

try:
    import warnings
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", FutureWarning)
        import pynvml
    pynvml.nvmlInit()
    _NVML = True
except Exception:
    _NVML = False


@dataclass
class GpuInfo:
    index: int
    name: str
    total_mb: float
    used_mb: float
    free_mb: float
    temperature: Optional[int] = None
    utilization_pct: Optional[int] = None
    mem_utilization_pct: Optional[int] = None
    power_w: Optional[float] = None
    fan_pct: Optional[int] = None


@dataclass
class ProcessGpuMem:
    pid: int
    dedicated_bytes: int = 0
    shared_bytes: int = 0
    name: str = ""
    exe: str = ""
    cmdline: str = ""
    gpu_index: int = 0


def get_gpu_info() -> list[GpuInfo]:
    gpus = _get_gpu_info_nvml()
    return gpus if gpus else _get_gpu_info_smi()


def _get_gpu_info_nvml() -> list[GpuInfo]:
    if not _NVML:
        return []
    gpus = []
    for i in range(pynvml.nvmlDeviceGetCount()):
        h = pynvml.nvmlDeviceGetHandleByIndex(i)
        mem = pynvml.nvmlDeviceGetMemoryInfo(h)
        name = pynvml.nvmlDeviceGetName(h)
        if isinstance(name, bytes):
            name = name.decode()

        temp = util_gpu = util_mem = power = fan = None
        try:
            temp = pynvml.nvmlDeviceGetTemperature(h, pynvml.NVML_TEMPERATURE_GPU)
        except Exception:
            pass
        try:
            u = pynvml.nvmlDeviceGetUtilizationRates(h)
            util_gpu, util_mem = u.gpu, u.memory
        except Exception:
            pass
        try:
            power = round(pynvml.nvmlDeviceGetPowerUsage(h) / 1000, 1)
        except Exception:
            pass
        try:
            fan = pynvml.nvmlDeviceGetFanSpeed(h)
        except Exception:
            pass

        gpus.append(GpuInfo(
            index=i, name=name,
            total_mb=mem.total / 1048576,
            used_mb=mem.used / 1048576,
            free_mb=mem.free / 1048576,
            temperature=temp,
            utilization_pct=util_gpu,
            mem_utilization_pct=util_mem,
            power_w=power,
            fan_pct=fan,
        ))
    return gpus


def _get_gpu_info_smi() -> list[GpuInfo]:
    try:
        r = subprocess.run(
            ["nvidia-smi",
             "--query-gpu=index,name,memory.total,memory.used,memory.free",
             "--format=csv,noheader,nounits"],
            capture_output=True, text=True, timeout=10,
        )
        gpus = []
        for line in r.stdout.strip().splitlines():
            parts = [p.strip() for p in line.split(",")]
            if len(parts) >= 5:
                gpus.append(GpuInfo(
                    index=int(parts[0]), name=parts[1],
                    total_mb=float(parts[2]),
                    used_mb=float(parts[3]),
                    free_mb=float(parts[4]),
                ))
        return gpus
    except Exception:
        return []


def _query_perf_counter(counter: str) -> dict[int, int]:
    """Query a Windows GPU Process Memory perf counter → {pid: bytes}."""
    try:
        r = subprocess.run(
            ["powershell", "-NoProfile", "-Command",
             f"Get-Counter '\\GPU Process Memory(*)\\{counter}' "
             "-ErrorAction SilentlyContinue "
             "| Select-Object -ExpandProperty CounterSamples "
             "| ForEach-Object { $_.InstanceName + '=' + $_.CookedValue }"],
            capture_output=True, text=True, timeout=15,
        )
    except Exception:
        return {}
    result: dict[int, int] = {}
    for line in r.stdout.strip().splitlines():
        m = re.match(r"pid_(\d+)_.*=(\d+(?:\.\d+)?)", line.strip())
        if m:
            pid = int(m.group(1))
            val = int(float(m.group(2)))
            result[pid] = result.get(pid, 0) + val
    return result


def get_process_gpu_memory() -> list[ProcessGpuMem]:
    dedicated = _query_perf_counter("Dedicated Usage")
    shared = _query_perf_counter("Shared Usage")
    all_pids = set(dedicated) | set(shared)
    procs: list[ProcessGpuMem] = []

    for pid in all_pids:
        ded = dedicated.get(pid, 0)
        shr = shared.get(pid, 0)
        if ded == 0 and shr == 0:
            continue

        name = exe = cmdline = ""
        if _psutil:
            try:
                p = _psutil.Process(pid)
                name = p.name()
                exe = p.exe()
                try:
                    cmdline = " ".join(p.cmdline())
                except (getattr(_psutil, 'AccessDenied', Exception),
                        getattr(_psutil, 'ZombieProcess', Exception)):
                    pass
            except Exception:
                pass

        procs.append(ProcessGpuMem(
            pid=pid, dedicated_bytes=ded, shared_bytes=shr,
            name=name, exe=exe, cmdline=cmdline,
        ))

    procs.sort(key=lambda p: p.dedicated_bytes, reverse=True)
    return procs


def short_name(proc: ProcessGpuMem) -> str:
    if proc.cmdline and "python" in proc.name.lower():
        for part in proc.cmdline.split():
            if part.endswith(".py"):
                return f"python {os.path.basename(part)}"
        return proc.name
    return proc.name or f"<pid {proc.pid}>"


def fmt_bytes(b: int) -> str:
    if b >= 1073741824:
        return f"{b / 1073741824:.2f} GB"
    if b >= 1048576:
        return f"{b / 1048576:.1f} MB"
    if b >= 1024:
        return f"{b / 1024:.0f} KB"
    return f"{b} B"


def kill_process(pid: int) -> dict:
    if not _psutil:
        return {"ok": False, "error": "psutil not installed"}
    try:
        p = _psutil.Process(pid)
        name = p.name()
        p.terminate()
        try:
            p.wait(timeout=3)
        except _psutil.TimeoutExpired:
            p.kill()
        return {"ok": True, "message": f"Terminated {name} (PID {pid})"}
    except _psutil.NoSuchProcess:
        return {"ok": False, "error": f"Process {pid} not found"}
    except _psutil.AccessDenied:
        return {"ok": False, "error": f"Access denied for PID {pid}"}
    except Exception as e:
        return {"ok": False, "error": str(e)}


def collect_snapshot() -> dict:
    gpus = get_gpu_info()
    procs = get_process_gpu_memory()
    total_gpu_bytes = int(gpus[0].total_mb * 1048576) if gpus else 1

    return {
        "gpus": [{
            **asdict(g),
            "used_pct": round(g.used_mb / g.total_mb * 100, 1) if g.total_mb else 0,
        } for g in gpus],
        "processes": [{
            "pid": p.pid,
            "name": short_name(p),
            "exe": p.exe,
            "dedicated_bytes": p.dedicated_bytes,
            "dedicated_human": fmt_bytes(p.dedicated_bytes),
            "shared_bytes": p.shared_bytes,
            "shared_human": fmt_bytes(p.shared_bytes),
            "gpu_pct": round(p.dedicated_bytes / total_gpu_bytes * 100, 1),
        } for p in procs],
        "total_tracked": fmt_bytes(sum(p.dedicated_bytes for p in procs)),
    }
