#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
GPU 即時數據抓取（方案 B 數據源）
讀取 NVIDIA GB10 的：温度 / 使用率 / 功率，並主機健康狀態：
CPU / 内存 / 磁盘 / 異常進程。輸出 JSON。
供同目錄 gpu_dashboard.html 定時拉取（每 3 秒）。

用法：
  python3 gpu_metric.py        # 輸出一筆 JSON 到 stdout
  python3 gpu_metric.py watch  # 每秒印出（除錯用）
"""
import subprocess
import json
import sys
import time
import os
import shutil

# 85°C 警報閾值（與 GPU 高溫警報排程一致）
ALARM_C = 85

# 主機健康閾值
CPU_WARN = 90   # CPU 使用率 %
MEM_WARN = 80   # 内存使用率 %
DISK_WARN = 85  # 磁盘使用率 %
PROC_WARN = 5   # 異常進程數上限（D=不可中斷 / Z=僵尸）


def read_system():
    """讀取主機健康狀態：CPU / 内存 / 磁盘 / 異常進程。"""
    try:
        # CPU 使用率 = 1 - idle（/proc/stat 第一行）
        with open('/proc/stat') as f:
            parts = f.readline().split()[1:]
        idle = int(parts[3]) + int(parts[4])
        total = sum(int(x) for x in parts)
        cpu = round((1 - idle/total) * 100, 1) if total else 0.0
    except Exception:
        cpu = None
    try:
        # 内存使用率
        with open('/proc/meminfo') as f:
            info = {l.split(':')[0]: int(l.split()[1]) for l in f}
        total = info.get('MemTotal', 0)
        avail = info.get('MemAvailable', 0)
        mem = round((1 - avail/total) * 100, 1) if total else None
    except Exception:
        mem = None
    try:
        # 磁盘使用率（/）
        disk = shutil.disk_usage('/')
        disk_pct = round(disk.used/disk.total*100, 1) if disk.total else None
    except Exception:
        disk_pct = None
    try:
        # 異常進程數（狀態 D / Z）
        anom = 0
        for p in os.listdir('/proc'):
            if not p.isdigit():
                continue
            try:
                with open(f'/proc/{p}/stat') as f:
                    st = f.read()
                state = st[st.rfind(')')+2:st.rfind(')')+3]
                if state in ('D', 'Z'):
                    anom += 1
            except Exception:
                continue
    except Exception:
        anom = None
    if cpu is None and mem is None and disk_pct is None and anom is None:
        return {"ok": False, "err": "host metrics failed"}
    healthy = (cpu is None or cpu < CPU_WARN) and \
              (mem is None or mem < MEM_WARN) and \
              (disk_pct is None or disk_pct < DISK_WARN) and \
              (anom is None or anom < PROC_WARN)
    return {
        "ok": True,
        "cpu_pct": cpu,
        "mem_pct": mem,
        "disk_pct": disk_pct,
        "anom_procs": anom,
        "healthy": healthy,
    }


def read_gpu():
    out = subprocess.run(
        ["nvidia-smi", "--query-gpu=temperature.gpu,utilization.gpu,power.draw",
         "--format=csv,noheader,nounits"],
        capture_output=True, text=True, timeout=15
    )
    if out.returncode != 0 or not out.stdout.strip():
        return {"ok": False, "err": out.stderr.strip() or "nvidia-smi failed"}
    parts = [p.strip() for p in out.stdout.strip().split(",")]
    try:
        temp = float(parts[0])
        util = float(parts[1])
        power = float(parts[2])
    except (ValueError, IndexError):
        return {"ok": False, "err": f"parse failed: {parts}", "raw": out.stdout.strip()}
    return {
        "ok": True,
        "temp_c": temp,
        "util_pct": util,
        "power_w": power,
        "alarm": temp >= ALARM_C,
        "alarm_c": ALARM_C,
        "host": read_system(),
        "host_name": os.uname().nodename,
        "ts": int(time.time() * 1000),
    }


if __name__ == "__main__":
    if "--watch" in sys.argv or "watch" in sys.argv:
        try:
            while True:
                print(json.dumps(read_gpu(), ensure_ascii=False))
                sys.stdout.flush()
                time.sleep(1)
        except KeyboardInterrupt:
            pass
    else:
        print(json.dumps(read_gpu(), ensure_ascii=False))