"""Lightweight perf sampler: GPU (nvidia-smi) + system/per-process CPU every 30s.

Writes CSV rows to logs/perf_samples.csv. Designed to run all day in background
so TTS synthesis duration (logs/tts_memory.log) can be correlated against
GPU temperature/clocks and CPU pressure afterwards.
"""

from __future__ import annotations

import csv
import shutil
import subprocess
import time
from datetime import datetime
from pathlib import Path

import psutil

ROOT = Path(__file__).resolve().parent.parent
OUT = ROOT / "logs" / "perf_samples.csv"
INTERVAL_SEC = 30
NVIDIA_SMI = shutil.which("nvidia-smi") or r"C:\Windows\System32\nvidia-smi.exe"

# process name substrings (lowercase) to track individually
TRACKED = {
    "livehime": "livehime.exe",
    "yuanshen": "yuanshen.exe",
    "bettergi": "bettergi.exe",
    "mpv": "mpv.exe",
    "audiodg": "audiodg.exe",
    "python": None,  # aggregated below (main + TTS worker)
}

FIELDNAMES = [
    "ts",
    "gpu_temp_c", "gpu_clock_mhz", "gpu_mem_clock_mhz", "gpu_util_pct",
    "gpu_power_w", "gpu_throttle",
    "cpu_total_pct",
    "cpu_livehime", "cpu_yuanshen", "cpu_bettergi", "cpu_mpv", "cpu_audiodg",
    "cpu_python_total", "python_proc_count",
]


def query_gpu() -> dict:
    keys = (
        "temperature.gpu", "clocks.current.graphics", "clocks.current.memory",
        "utilization.gpu", "power.draw", "clocks_event_reasons.active",
    )
    try:
        out = subprocess.check_output(
            [NVIDIA_SMI, f"--query-gpu={','.join(keys)}", "--format=csv,noheader,nounits"],
            text=True, timeout=5,
            creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0),
        ).strip()
        temp, clk, memclk, util, power, throttle = [x.strip() for x in out.split(",")]
        return {
            "gpu_temp_c": temp,
            "gpu_clock_mhz": clk,
            "gpu_mem_clock_mhz": memclk,
            "gpu_util_pct": util,
            "gpu_power_w": power,
            "gpu_throttle": throttle,
        }
    except Exception:
        return {k: "" for k in FIELDNAMES if k.startswith("gpu_")}


def sample_cpu(procs: dict[str, list[psutil.Process]]) -> dict:
    row = {"cpu_total_pct": psutil.cpu_percent(interval=None)}
    py_total = 0.0
    py_count = 0
    for name, plist in procs.items():
        total = 0.0
        alive = []
        for p in plist:
            try:
                total += p.cpu_percent(interval=None)
                alive.append(p)
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                pass
        procs[name] = alive
        if name == "python":
            py_total, py_count = total, len(alive)
        else:
            row[f"cpu_{name}"] = round(total, 1)
    row["cpu_python_total"] = round(py_total, 1)
    row["python_proc_count"] = py_count
    return row


def find_tracked() -> dict[str, list[psutil.Process]]:
    found: dict[str, list[psutil.Process]] = {k: [] for k in TRACKED}
    for p in psutil.process_iter(["name"]):
        try:
            pname = (p.info["name"] or "").lower()
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            continue
        for key in TRACKED:
            if key in pname:
                found[key].append(p)
    return found


def main() -> None:
    OUT.parent.mkdir(parents=True, exist_ok=True)
    new_file = not OUT.exists()
    procs = find_tracked()
    # prime cpu_percent baselines
    psutil.cpu_percent(interval=None)
    for plist in procs.values():
        for p in plist:
            try:
                p.cpu_percent(interval=None)
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                pass
    with OUT.open("a", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=FIELDNAMES)
        if new_file:
            writer.writeheader()
        while True:
            time.sleep(INTERVAL_SEC)
            if time.monotonic() % 300 < INTERVAL_SEC:  # re-discover processes ~every 5 min
                procs = find_tracked()
            row = {"ts": datetime.now().isoformat(timespec="seconds")}
            row.update(query_gpu())
            row.update(sample_cpu(procs))
            writer.writerow(row)
            f.flush()


if __name__ == "__main__":
    main()