#!/usr/bin/env python3
"""Radiator fan control for lovey: fans follow max(CPU curve, GPU curve).

CPU temp comes from coretemp. The GPUs are passed through to VM 200, so the
guest pushes its hottest GPU temp here over UDP every 2 s (gpu-temp-push in
the guest's NixOS config). On exit the fans go back to BIOS control.
"""
import glob, os, signal, socket, sys, time

FANS = [2, 3, 4, 5]            # nct6798 headers with radiator fans; fan6 (AIO_PUMP) and fan7 (pump) stay on BIOS
LISTEN = ("192.168.1.53", 7777)
GPU_SOURCE = "192.168.1.67"    # llm-nixos; only accept temps from it
VMID = 200
INTERVAL = 2.0
STALE = 15.0                   # seconds without a GPU packet while the VM runs -> failsafe
FAILSAFE = 85                  # % when GPU temp is unknown but the VM is up
MIN_PCT = 30
RAMP_DOWN = 2                  # max % drop per interval (100 -> 30 in ~70 s)
ALPHA = 0.3                    # temperature smoothing per interval (~6 s time constant)

# (temp C, fan %) points, linearly interpolated, flat outside the range
CPU_CURVE = [(50, 30), (70, 50), (85, 100)]
GPU_CURVE = [(45, 30), (60, 40), (75, 100)]


def hwmon(name):
    for d in glob.glob("/sys/class/hwmon/hwmon*"):
        with open(f"{d}/name") as f:
            if f.read().strip() == name:
                return d
    sys.exit(f"hwmon {name} not found")


def read_int(path):
    with open(path) as f:
        return int(f.read())


def write(path, value):
    with open(path, "w") as f:
        f.write(str(value))


def curve(points, t):
    if t <= points[0][0]:
        return points[0][1]
    for (t0, p0), (t1, p1) in zip(points, points[1:]):
        if t <= t1:
            return p0 + (p1 - p0) * (t - t0) / (t1 - t0)
    return points[-1][1]


def vm_running():
    try:
        with open(f"/var/run/qemu-server/{VMID}.pid") as f:
            os.kill(int(f.read()), 0)
        return True
    except (OSError, ValueError):
        return False


NCT = hwmon("nct6798")
CORETEMP = hwmon("coretemp")
CPU_TEMP = next(p.replace("_label", "_input") for p in glob.glob(f"{CORETEMP}/temp*_label")
                if open(p).read().startswith("Package"))


def restore(*_):
    for n in FANS:
        try:
            write(f"{NCT}/pwm{n}_enable", 5)  # 5 = SmartFan IV, the BIOS curve
        except OSError:
            pass
    sys.exit(0)


def main():
    signal.signal(signal.SIGTERM, restore)
    signal.signal(signal.SIGINT, restore)
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.bind(LISTEN)
    sock.setblocking(False)
    for n in FANS:
        write(f"{NCT}/pwm{n}_enable", 1)

    cpu_s = gpu_s = None
    gpu_at = 0.0
    pct = 100.0
    last_log = 0.0
    try:
        while True:
            while True:  # drain queued GPU packets, keep the newest
                try:
                    data, (addr, _) = sock.recvfrom(64)
                except BlockingIOError:
                    break
                if addr == GPU_SOURCE:
                    try:
                        g = float(data.decode().strip())
                    except ValueError:
                        continue
                    gpu_s = g if gpu_s is None else gpu_s + ALPHA * (g - gpu_s)
                    gpu_at = time.monotonic()

            c = read_int(CPU_TEMP) / 1000
            cpu_s = c if cpu_s is None else cpu_s + ALPHA * (c - cpu_s)

            want = curve(CPU_CURVE, cpu_s)
            if time.monotonic() - gpu_at < STALE:
                want = max(want, curve(GPU_CURVE, gpu_s))
                gpu_state = f"{gpu_s:.0f}C"
            elif vm_running():
                want = max(want, FAILSAFE)
                gpu_state = "stale"
            else:
                gpu_state = "vm-off"
            want = max(want, MIN_PCT)
            pct = want if want > pct else max(want, pct - RAMP_DOWN)

            for n in FANS:
                write(f"{NCT}/pwm{n}", round(pct * 255 / 100))

            if time.monotonic() - last_log >= 60:
                rpm = " ".join(str(read_int(f"{NCT}/fan{n}_input")) for n in FANS)
                print(f"cpu={cpu_s:.0f}C gpu={gpu_state} fans={pct:.0f}% rpm={rpm}", flush=True)
                last_log = time.monotonic()
            time.sleep(INTERVAL)
    finally:
        restore()


if __name__ == "__main__":
    main()