#!/usr/bin/python3
import os
import re
import time

from ubnthelpers import get_ubnt_shortname

SMARTCTL_PATH = "/usr/sbin/smartctl"

THERMAL_SYS_PATHS = {
    "UNVRPRO": {
        "thermal": [
            "/sys/devices/virtual/thermal/thermal_zone0/temp",
            "/sys/class/hwmon/hwmon0/device/temp1_input",
            "/sys/class/hwmon/hwmon0/device/temp2_input",
            "/sys/class/hwmon/hwmon0/device/temp3_input",
        ],
        "fan_pwms": [
            "/sys/class/hwmon/hwmon0/device/pwm1",
            "/sys/class/hwmon/hwmon0/device/pwm2",
            "/sys/class/hwmon/hwmon0/device/pwm3",
        ],
    },
}


def __get_board_temps():
    # Are we supported?
    if get_ubnt_shortname() not in THERMAL_SYS_PATHS:
        raise Exception(
            f"Error: Your Unifi device of {get_ubnt_shortname()} is not yet supported by unvr-fan-daemon! Exiting..."
        )
    # For each of our paths, get the temps, and append to single return list
    board_temps = []
    for path in THERMAL_SYS_PATHS[get_ubnt_shortname()]["thermal"]:
        try:
            with open(path, "r") as f:
                board_temps.append(int(f.readline().rstrip("\n")))
        except FileNotFoundError:
            print(f"Warning: Unable to open {path}; ignoring and continuing...")
            continue

    # Did we get ANY temps?!?
    if len(board_temps) == 0:
        raise Exception(
            "Error: Unable to parse out any board temps for your device, something is really wrong! Exiting..."
        )

    return board_temps


def __get_disk_temps():
    # Find the list of all devices, which could be none
    devices = re.findall(
        r"^[/a-z]+",
        os.popen(f"{SMARTCTL_PATH} -n standby --scan").read(),
        re.MULTILINE,
    )

    # For each disk, get the temp, and append to our list
    disk_temps = []
    for dev in devices:
        # Sadly this is slow, SMART data pulls are not fast...
        dev_temp = re.search(
            r"^194 [\w-]+\s+0x\d+\s+\d+\s+\d+\s+\d+\s+[\w-]+\s+\w+\s+\S+\s+(\d+)(?:\s[\(][^)]*[\)])?$",
            os.popen(f"{SMARTCTL_PATH} -A {dev}").read(),
            re.MULTILINE,
        )
        if dev_temp:
            disk_temps.append(int(f"{dev_temp.group(1)}000"))  # Append zeros

    return disk_temps


def __calculate_fan_speed(temp):
    # our basic fancurve logic
    match temp:
        case _ if temp < 40:
            fanspeed = 25
        case _ if temp >= 40 and temp < 50:
            fanspeed = 75
        case _ if temp >= 50 and temp < 55:
            fanspeed = 100
        case _ if temp >= 55 and temp < 60:
            fanspeed = 130
        case _ if temp >= 60 and temp < 65:
            fanspeed = 160
        case _ if temp >= 65 and temp < 70:
            fanspeed = 200
        case _:
            fanspeed = 255

    return fanspeed


def __set_fan_speed(speed: int):
    # Set the fans
    for fan in THERMAL_SYS_PATHS[get_ubnt_shortname()]["fan_pwms"]:
        try:
            with open(fan, "w") as f:
                f.write(str(speed))
        except FileNotFoundError:
            print(
                f"Error: Unable to write to PWM at {fan}! Why can't we set fan speed!?"
            )
            raise


def __write_out_temp(temp: int, path: str = "/tmp/.unvr_temp"):
    try:
        with open(path, "w+") as f:
            f.write(str(temp))
    except (IOError, OSError, PermissionError) as e:
        print(
            f"Warning: Unable to write to temp file at {path}; ulcmd won't get the system temp! Error was: {e}"
        )


if __name__ == "__main__":
    # Cache so we only write to PWMs if this changes
    last_fanspeed = 0

    print("unvr-fan-daemon starting...")

    # Start with debug write to max speed so we hear it :)
    __set_fan_speed(255)
    time.sleep(0.5)

    # Start our main loop
    while True:
        # Get the fanspeed we wanna set based on temps
        temp = (
            sorted(__get_board_temps() + __get_disk_temps(), reverse=True)[0] / 1000.0
        )  # Move temp to C

        fanspeed = __calculate_fan_speed(temp)

        # Write out for consumption by ulcmd via mock-ubnt-api
        __write_out_temp(temp)

        # If there's a change in calculated fan speed, set it
        if last_fanspeed != fanspeed:
            print(f"Setting fan PWMs to {fanspeed} due to temp of {temp}C")
            __set_fan_speed(fanspeed)
            last_fanspeed = fanspeed

        # Sleep and run again
        time.sleep(10)