"""Low-frequency NVIDIA-SMI textfile collector; no GPU context or driver changes."""
import csv
import io
import math
import os
from pathlib import Path
import signal
import subprocess
import sys
import time

running = True


def stop(*_):
    global running
    running = False


def sample():
    rows = ["# TYPE contextra_gpu_collector_success gauge"]
    try:
        value = subprocess.run(["nvidia-smi", "--query-gpu=index,utilization.gpu,memory.used,memory.total,temperature.gpu,power.draw",
            "--format=csv,noheader,nounits"], check=True, capture_output=True, text=True, timeout=5)
        names = ("utilization_percent", "memory_used_mib", "memory_total_mib", "temperature_celsius", "power_watts")
        for row in csv.reader(io.StringIO(value.stdout)):
            index = int(row[0])
            for name, raw in zip(names, row[1:]):
                try:
                    number = float(raw.strip())
                    if math.isfinite(number):
                        rows.append(f'contextra_gpu_{name}{{gpu="{index}"}} {number}')
                except ValueError:
                    pass
        rows.append("contextra_gpu_collector_success 1")
    except (subprocess.SubprocessError, OSError, ValueError):
        rows.append("contextra_gpu_collector_success 0")
    rows.append(f"contextra_gpu_collector_timestamp_seconds {time.time()}")
    return "\n".join(rows) + "\n"


def main():
    target = Path(sys.argv[1])
    if not target.is_dir() or target.is_symlink():
        raise RuntimeError("Metrics directory missing")
    signal.signal(signal.SIGTERM, stop)
    signal.signal(signal.SIGINT, stop)
    while running:
        temporary = target / f".gpu-{os.getpid()}.tmp"
        temporary.write_text(sample())
        temporary.chmod(0o644)
        temporary.replace(target / "gpu.prom")
        for _ in range(15):
            if not running:
                break
            time.sleep(1)


if __name__ == "__main__":
    main()