import shutil
import subprocess
import time
from typing import Any

import psutil

_prev_io: dict[int, tuple[float, float, float]] = {}
_gpu_procs: dict[int, int] | None = None
_gpu_probe_t = 0.0


def _gpu_per_proc() -> dict[int, int]:
    """Map PID to GPU memory used (MiB) for NVIDIA compute processes.

    Runs `nvidia-smi --query-compute-apps` at most once per 10 seconds
    (the probe result is cached). Returns an empty mapping when nvidia-smi
    is missing, which is the case on AMD machines.

    Returns:
        A pid to used-memory-in-MiB mapping.
    """
    global _gpu_procs, _gpu_probe_t
    if not shutil.which("nvidia-smi"):
        return {}
    if _gpu_procs is not None and time.monotonic() - _gpu_probe_t < 10:
        return _gpu_procs
    _gpu_probe_t = time.monotonic()
    _gpu_procs = {}
    try:
        out = subprocess.run(
            [
                "nvidia-smi",
                "--query-compute-apps=pid,used_memory",
                "--format=csv,noheader,nounits",
            ],
            capture_output=True,
            text=True,
            timeout=5,
            check=False
        ).stdout
        for line in out.splitlines():
            parts = [p.strip() for p in line.split(",")]
            if len(parts) >= 2:
                try:
                    _gpu_procs[int(parts[0])] = int(parts[1])
                except ValueError:
                    continue
    except (OSError, subprocess.SubprocessError):
        pass
    return _gpu_procs


def sample() -> list[dict[str, Any]]:
    """One pass over all processes collecting cpu, memory, IO rate, GPU.

    Processes whose parent is swapper/kthreadd (ppid 0/2) are skipped.
    Per-process IO rates are byte deltas between successive calls divided
    by elapsed time; previous readings are pruned when a process exits.
    GPU memory comes from _gpu_per_proc(). Entries that die mid-iteration
    are dropped, and per-process access errors are tolerated.

    Returns:
        A list of per-process dicts (pid, name, user, cpu, rss, mem_pct,
        io_read, io_write, gpu), one per live process.
    """
    now = time.monotonic()
    mem_total = psutil.virtual_memory().total
    gpu = _gpu_per_proc()
    out: list[dict[str, Any]] = []
    alive: set[int] = set()
    for p in psutil.process_iter():
        try:
            with p.oneshot():
                if p.ppid() in (0, 2):
                    continue
                cpu = p.cpu_percent(None)
                mem = p.memory_info()
                name = p.name()
                user = p.username()
                try:
                    io = p.io_counters()
                except (psutil.AccessDenied, psutil.NoSuchProcess, OSError):
                    io = None
        except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
            continue
        pid = p.pid
        alive.add(pid)
        entry = {
            "pid": pid,
            "name": name,
            "user": user,
            "cpu": cpu,
            "rss": mem.rss,
            "mem_pct": (mem.rss / mem_total * 100) if mem_total else 0.0,
            "io_read": 0.0,
            "io_write": 0.0,
            "gpu": gpu.get(pid),
        }
        if io is not None:
            prev = _prev_io.get(pid)
            if prev is not None and now > prev[0]:
                dt = now - prev[0]
                entry["io_read"] = max(0.0, (io.read_bytes - prev[1]) / dt)
                entry["io_write"] = max(0.0, (io.write_bytes - prev[2]) / dt)
            _prev_io[pid] = (now, io.read_bytes, io.write_bytes)
        out.append(entry)
    for pid in list(_prev_io):
        if pid not in alive:
            del _prev_io[pid]
    return out