"""Query GPU stats via nvidia-smi."""

from __future__ import annotations

import json
import shutil
import subprocess
from typing import Any


class GpuInfoError(RuntimeError):
    pass


def query_gpu() -> dict[str, Any]:
    if not shutil.which("nvidia-smi"):
        return {
            "available": False,
            "error": "nvidia-smi not found",
            "gpus": [],
        }

    try:
        result = subprocess.run(
            [
                "nvidia-smi",
                "--query-gpu=index,name,memory.used,memory.total,temperature.gpu,utilization.gpu",
                "--format=csv,noheader,nounits",
            ],
            capture_output=True,
            text=True,
            timeout=15,
            check=False,
        )
    except subprocess.TimeoutExpired as exc:
        raise GpuInfoError("nvidia-smi timed out") from exc

    if result.returncode != 0:
        return {
            "available": False,
            "error": (result.stderr or result.stdout or "nvidia-smi failed").strip(),
            "gpus": [],
        }

    gpus: list[dict[str, Any]] = []
    for line in result.stdout.strip().splitlines():
        parts = [p.strip() for p in line.split(",")]
        if len(parts) < 6:
            continue
        idx, name, mem_used, mem_total, temp, util = parts[:6]
        try:
            used_mb = int(float(mem_used))
            total_mb = int(float(mem_total))
        except ValueError:
            used_mb = 0
            total_mb = 0
        gpus.append(
            {
                "index": int(idx),
                "name": name,
                "memory_used_mb": used_mb,
                "memory_total_mb": total_mb,
                "memory_used_pct": round(100 * used_mb / total_mb, 1) if total_mb else 0,
                "temperature_c": int(float(temp)) if temp not in ("N/A", "") else None,
                "utilization_pct": int(float(util)) if util not in ("N/A", "") else None,
            }
        )

    return {"available": True, "error": None, "gpus": gpus}