#!python
from cm_rgb.ctrl import LedChannel, LedMode, CMRGBController
import psutil
import atexit
import time
import click

# Try to load optional dependency - sensors
sensorsImported = True
sensorsImportError = ""

try:
    import sensors
    atexit.register(sensors.cleanup)
    sensors.init()

except ImportError as error:
    sensorsImported = False
    sensorsImportError = error


def verify_sensors_import():
    # Make sure optional dependency is available for showSensor
    if not sensorsImported:
        print(sensorsImportError,
          "\n\nCould not import sensors.\n" +
          "Cannot show temperature as fan color.\n\n" +
          "To fix this, run:\npip3 install pysensors"
          )
        exit(-1)

def print_available_sources(ctx,value):
    if value:
        verify_sensors_import()

        print("Available sensors:")
        for chip in sensors.iter_detected_chips():
            for feature in chip:
                print(str(chip)+"/"+feature.label," -> ",feature.get_value())

        ctx.exit()


@click.command()
@click.option("--bg-color", "bgColor", default="#00FFFF", help="Background LED's color")
@click.option("--cpu-color", "cpuColor", default="#FFA500", help="Color of the cpu load LED's")
@click.option("--brightness", type=click.IntRange(1, 5, clamp=True), default=3)
@click.option("--interval", type=click.FloatRange(0.01, 60, clamp=True), default=0.2)

@click.option("--show-temp", "showSensor", is_flag=True, help="Show temperature of selected sensor on cpu fan")
@click.option("--temp-source", "tempSource", type=str, default="k10temp-pci-00c3/Tdie", help="Temperature source <chip>/<feature> (eg. \"k10temp-pci-00c3/Tdie\")")

@click.option("--temp-low", "low", type=float, default=45, help="Temperature considered low")
@click.option("--temp-high", "high", type=float, default=85, help="Temperature considered high")

@click.option("--temp-low-color", "lowColor", default="#00FFFF", help="Color representing low temperature")
@click.option("--temp-high-color", "highColor", default="#FFA500", help="Color representing high temperature")

@click.option("--print-values", "printValues", is_flag=True, help="Print cpu load and sensor readout")

@click.option('--list-temp-sources', is_flag=True, callback=print_available_sources, expose_value=False, is_eager=True)

def monitor(
        bgColor,
        cpuColor,
        brightness,
        interval,
        showSensor,
        low,
        high,
        lowColor,
        highColor,
        tempSource,
        printValues
        ):
    c = CMRGBController()

    bgChannel = LedChannel.R_STATIC
    cpuChannel = LedChannel.R_SWIRL

    b = [0x33, 0x66, 0x99, 0xCC, 0xFF][brightness-1]

    bgColor = bgColor.lstrip('#')
    col = [int(bgColor[i:i+2], 16) for i in (0, 2, 4)]
    c.set_channel(bgChannel, LedMode.R_DEFAULT, b, col[0], col[1], col[2])

    cpuColor = cpuColor.lstrip('#')
    col = [int(cpuColor[i:i+2], 16) for i in (0, 2, 4)]
    c.set_channel(cpuChannel, LedMode.R_DEFAULT, b, col[0], col[1], col[2], 0x60)

    c.apply()

    def exit_handler():
        c.restore()

    atexit.register(exit_handler)

    lowColor = [int(lowColor.lstrip('#')[i:i+2], 16) for i in (0, 2, 4)]
    highColor = [int(highColor.lstrip('#')[i:i+2], 16) for i in (0, 2, 4)]

    if showSensor:
        verify_sensors_import()
        sensor = False

        tmp = tempSource.split('/')
        if len(tmp) == 2:
            sensorChip = tmp[0]
            sensorFeature = tmp[1]

            for chip in sensors.iter_detected_chips(sensorChip):
                
                fGen = (f for f in chip if (f.label == sensorFeature))
                for feature in fGen:
                    sensor = feature
        else:
            print("Invalid temp source syntax. Must match <chip>/<feature>")
            exit(-1)

        if not sensor:
            print("Temp source not found, try running with --list-temp-sources")
            exit(-1)

    while True:
        if showSensor:
            # print(sensor.label, sensor.get_value())
            t = sensor.get_value()
            interp_t = max(0, min(1, (t-low)/(high-low)))
            color = [
                int(
                    interp_t * highColor[i]
                    + (1 - interp_t)*lowColor[i]
                    )
                for i in range(3)
                ]
            # print(color)
            c.set_channel(LedChannel.FAN, LedMode.STATIC, b, color[0], color[1], color[2])

        # gives a single float value
        cpu = psutil.cpu_percent()
        cpu_leds = int(round(cpu*15 / 100))

        total = 15 - cpu_leds

        ring_leds = ([cpuChannel]*cpu_leds)
        ring_leds = ring_leds + ([bgChannel]*total)

        shift = -8
        ring_leds = ring_leds[-shift:]+ring_leds[:-shift]

        c.assign_leds_to_channels(LedChannel.LOGO, LedChannel.FAN, *ring_leds)
        c.apply()

        if printValues:
            if showSensor:
                print("CPU:",cpu,"TEMP:",t)
            else:
                print("CPU:",cpu)
        time.sleep(interval)


if __name__ == '__main__':
    monitor()
