# GPU监控模块 · 采集nvidia-smi数据
# HLDP://zhuyuan-agent/gpu-monitor

import subprocess
import json
import re
from typing import List, Dict, Optional


def collect_gpu_metrics() -> Dict:
    """采集所有GPU的实时指标
    
    Returns:
        {
            "gpus": [
                {
                    "index": 0,
                    "name": "NVIDIA GeForce RTX 3090",
                    "uuid": "GPU-xxxx",
                    "utilization_gpu": 85,     # %
                    "memory_used_mb": 18432,   # MB
                    "memory_total_mb": 24576,  # MB
                    "temperature_gpu": 72,     # °C
                    "power_draw_w": 285.5,     # W
                    "fan_speed": 65            # %
                }
            ],
            "error": null
        }
    """
    try:
        # 查询GPU关键指标
        result = subprocess.run(
            [
                "nvidia-smi",
                "--query-gpu=index,name,uuid,utilization.gpu,memory.used,memory.total,temperature.gpu,power.draw,fan.speed",
                "--format=csv,noheader,nounits"
            ],
            capture_output=True, text=True, timeout=10
        )
        
        if result.returncode != 0:
            return {"gpus": [], "error": f"nvidia-smi failed: {result.stderr.strip()}"}
        
        gpus = []
        for line in result.stdout.strip().split("\n"):
            if not line.strip():
                continue
            parts = [p.strip() for p in line.split(",")]
            if len(parts) < 9:
                continue
            
            try:
                gpu = {
                    "index": int(parts[0]),
                    "name": parts[1],
                    "uuid": parts[2],
                    "utilization_gpu": int(parts[3]) if parts[3] != "[Not Supported]" else 0,
                    "memory_used_mb": int(parts[4]) if parts[4] != "[Not Supported]" else 0,
                    "memory_total_mb": int(parts[5]) if parts[5] != "[Not Supported]" else 0,
                    "temperature_gpu": int(parts[6]) if parts[6] != "[Not Supported]" else 0,
                    "power_draw_w": float(parts[7]) if parts[7] not in ("[Not Supported]", "") else 0.0,
                    "fan_speed": int(parts[8]) if parts[8] not in ("[Not Supported]", "") else 0,
                }
                gpus.append(gpu)
            except (ValueError, IndexError):
                continue
        
        return {"gpus": gpus, "error": None}
    
    except FileNotFoundError:
        return {"gpus": [], "error": "nvidia-smi not found - not a GPU machine?"}
    except subprocess.TimeoutExpired:
        return {"gpus": [], "error": "nvidia-smi timed out"}
    except Exception as e:
        return {"gpus": [], "error": str(e)}


def gpu_summary(gpus: List[Dict]) -> str:
    """生成GPU状态的一行摘要"""
    if not gpus:
        return "无GPU"
    parts = []
    for g in gpus:
        util = g.get("utilization_gpu", 0)
        temp = g.get("temperature_gpu", 0)
        mem = g.get("memory_used_mb", 0)
        mem_total = g.get("memory_total_mb", 0)
        mem_pct = int(mem / mem_total * 100) if mem_total > 0 else 0
        parts.append(f"GPU{g['index']}: {util}%/{temp}°C/{mem_pct}%VRAM")
    return " | ".join(parts)


# 快速测试
if __name__ == "__main__":
    data = collect_gpu_metrics()
    print(json.dumps(data, indent=2, ensure_ascii=False))
    if data["gpus"]:
        print(f"\n摘要: {gpu_summary(data['gpus'])}")