#!/usr/bin/env python3
"""
NVIDIA Aggressive Fan Control Daemon
RTX Pro 6000 Blackwell - Gaming-style fan curve
Runs headless, no X11 required.
"""

import pynvml
import time
import logging
import sys

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s %(levelname)s %(message)s',
    handlers=[logging.StreamHandler(sys.stdout)]
)
log = logging.getLogger(__name__)

# --- Fan curve: (temp_C, fan_pct) ---
# More aggressive than stock workstation defaults.
# By 70C fans are already at 80%, by 80C they're maxed.
FAN_CURVE = [
    (0,  30),
    (45, 40),
    (55, 55),
    (63, 70),
    (70, 82),
    (75, 90),
    (80, 100),
]

POLL_INTERVAL = 3   # seconds between checks
GPU_INDEX     = 0   # change if you have multiple GPUs


def interpolate_fan(temp):
    """Return target fan % for a given GPU temp using the curve above."""
    if temp <= FAN_CURVE[0][0]:
        return FAN_CURVE[0][1]
    if temp >= FAN_CURVE[-1][0]:
        return FAN_CURVE[-1][1]
    for i in range(1, len(FAN_CURVE)):
        t1, f1 = FAN_CURVE[i - 1]
        t2, f2 = FAN_CURVE[i]
        if t1 <= temp <= t2:
            ratio = (temp - t1) / (t2 - t1)
            return int(f1 + ratio * (f2 - f1))
    return 100


def set_fans(handle, num_fans, pct):
    pct = max(0, min(100, pct))
    for fan_idx in range(num_fans):
        pynvml.nvmlDeviceSetFanSpeed_v2(handle, fan_idx, pct)


def main():
    pynvml.nvmlInit()
    driver = pynvml.nvmlSystemGetDriverVersion()
    log.info(f"NVIDIA driver: {driver}")

    handle   = pynvml.nvmlDeviceGetHandleByIndex(GPU_INDEX)
    name     = pynvml.nvmlDeviceGetName(handle)
    num_fans = pynvml.nvmlDeviceGetNumFans(handle)
    log.info(f"Controlling: {name}  ({num_fans} fans)")
    log.info("Fan curve (temp->fan%): " +
             "  ".join(f"{t}C->{f}%" for t, f in FAN_CURVE))

    # Take manual control of the fan
    for fan_idx in range(num_fans):
        pynvml.nvmlDeviceSetFanControlPolicy(
            handle, fan_idx,
            pynvml.NVML_FAN_POLICY_MANUAL
        )
    log.info("Fan control set to MANUAL")

    last_pct = -1
    try:
        while True:
            temp    = pynvml.nvmlDeviceGetTemperature(
                handle, pynvml.NVML_TEMPERATURE_GPU)
            target  = interpolate_fan(temp)

            if target != last_pct:
                set_fans(handle, num_fans, target)
                log.info(f"GPU {temp}C -> fans {target}%")
                last_pct = target

            time.sleep(POLL_INTERVAL)

    except KeyboardInterrupt:
        log.info("Interrupted - restoring automatic fan control")
    finally:
        # Hand control back to the driver on exit
        for fan_idx in range(num_fans):
            pynvml.nvmlDeviceSetFanControlPolicy(
                handle, fan_idx,
                pynvml.NVML_FAN_POLICY_TEMPERATURE_CONTINOUS_SW
            )
        pynvml.nvmlShutdown()
        log.info("Fan control restored to automatic")


if __name__ == "__main__":
    main()