#!/usr/bin/env python3
"""
GPU stats HTTP server for AMD ROCm (rocm-smi).
No external Python dependencies — uses rocm-smi subprocess for data.

Usage (host):
  python3 /opt/rocm/libexec/rocm_smi/../../.. && python3 gpu_stats.py
  # or just: python3 gpu_stats.py  (if /opt/rocm/bin is in PATH or ROCM_SMI set)

Usage (Docker — see docker-compose.yml gpu-stats service):
  Requires /opt/rocm mounted and /dev/kfd + /dev/dri devices passed through.

Dashboard config (config.json):
  "gpuStats": { "api": "/api/gpu-stats" }

Then uncomment the /api/gpu-stats block in nginx.conf.
"""
import json
import os
import subprocess
from http.server import BaseHTTPRequestHandler, HTTPServer

PORT     = 9101
ROCM_SMI = os.environ.get("ROCM_SMI_PATH", "/opt/rocm/bin/rocm-smi")


def _run(extra_flags):
    result = subprocess.run(
        [ROCM_SMI] + extra_flags + ["--json"],
        capture_output=True, text=True, timeout=5,
    )
    if result.returncode != 0:
        return {}
    try:
        return json.loads(result.stdout).get("card0", {})
    except Exception:
        return {}


def query_gpu():
    try:
        stats = _run(["--showuse", "--showtemp", "--showpower", "--showid"])
        vram_raw = subprocess.run(
            [ROCM_SMI, "--showmeminfo", "vram", "--json"],
            capture_output=True, text=True, timeout=5,
        )
        vram = {}
        if vram_raw.returncode == 0:
            vram = json.loads(vram_raw.stdout).get("card0", {})

        if not stats:
            return None

        vram_total_b = int(vram.get("VRAM Total Memory (B)", 0))
        vram_used_b  = int(vram.get("VRAM Total Used Memory (B)", 0))

        raw_name = stats.get("Card Series", stats.get("Device Name", ""))
        name     = raw_name if raw_name and raw_name != "N/A" else os.environ.get("GPU_NAME", "AMD GPU")
        return {
            "name":        name,
            "utilization": int(float(stats.get("GPU use (%)", 0))),
            "memory_used":  vram_used_b  // (1024 * 1024),   # MB
            "memory_total": vram_total_b // (1024 * 1024),   # MB
            "temperature": int(float(
                stats.get("Temperature (Sensor junction) (C)",
                stats.get("Temperature (Sensor edge) (C)", 0))
            )),
            "power_w":     float(stats.get("Average Graphics Package Power (W)", 0)),
            "power_max_w": float(stats.get("Max Graphics Package Power (W)", 0)),
        }
    except Exception:
        return None


class Handler(BaseHTTPRequestHandler):
    def do_GET(self):
        if self.path not in ("/", "/stats"):
            self.send_response(404)
            self.end_headers()
            return
        data = query_gpu()
        if data is None:
            body = json.dumps({"error": "rocm-smi unavailable"}).encode()
            self.send_response(503)
        else:
            body = json.dumps(data).encode()
            self.send_response(200)
        self.send_header("Content-Type", "application/json")
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def log_message(self, format, *args):  # noqa: A002
        pass  # suppress per-request noise


if __name__ == "__main__":
    server = HTTPServer(("0.0.0.0", PORT), Handler)
    print(f"GPU stats server (AMD ROCm) listening on :{PORT}")
    server.serve_forever()