#!/usr/bin/env python3
from cm_rgb.ctrl import LedChannel, LedMode, CMRGBController
import psutil
import atexit
import time
import click

def print_available_sources(ctx,value):
    if value:
        print("Available sensors:")
        sensorDict = psutil.sensors_temperatures()
        for chipname, chip in sensorDict.items():
            for feature in chip:
                print(chipname + "/" + feature.label," -> ", feature.current)

        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("--verbose", "verbose", is_flag=True, help="Print cpu load and sensor readout")

@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/Tdie", help="Temperature source <chip>/<feature> (eg. \"k10temp/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", "tLowColor", default="#00FFFF", help="Color representing low temperature")
@click.option("--temp-high-color", "tHighColor", default="#FFA500", help="Color representing high temperature")

@click.option("--show-cpu-frequency", "showCPUFreq", is_flag=True, help="Show CPU frequency on logo")
@click.option("--freq-low-color", "fLowColor", default="#00FFFF", help="Color representing low frequency")
@click.option("--freq-high-color", "fHighColor", default="#FFA500", help="Color representing high frequency")
@click.option("--freq-color-smoothing", "fSmoothing", type=click.FloatRange(0, 1, clamp=True), default=0.8, help="Smoothing of frequency value to make color less jumpy. 0 -> no smoothing, 0.9 -> a lot")

@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,
        tLowColor,
        tHighColor,
        tempSource,
        verbose,
        showCPUFreq,
        fLowColor,
        fHighColor,
        fSmoothing
        ):
    c = CMRGBController()

    c.disableMirage()

    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)


    tLowColor = [int(tLowColor.lstrip('#')[i:i+2], 16) for i in (0, 2, 4)]
    tHighColor = [int(tHighColor.lstrip('#')[i:i+2], 16) for i in (0, 2, 4)]
    
    fLowColor = [int(fLowColor.lstrip('#')[i:i+2], 16) for i in (0, 2, 4)]
    fHighColor = [int(fHighColor.lstrip('#')[i:i+2], 16) for i in (0, 2, 4)]
    
    smoothed_freq = 3000
    
    try:
        sensorGroup, sensorLabel = tempSource.split("/")
        if verbose:
            print(sensorGroup, sensorLabel)
        for i, sensor in enumerate(psutil.sensors_temperatures()[sensorGroup]):
            if sensor.label == sensorLabel:
                sensorIndex = i
    except:
        print("Temp source not found, try running with --list-temp-sources")
        showSensor = False

    while True:
        if showSensor:
            t = psutil.sensors_temperatures()[sensorGroup][sensorIndex].current
            interp_t = max(0, min(1, (t-low)/(high-low)))
            color = [
                int(interp_t * tHighColor[i] + (1 - interp_t)*tLowColor[i])
                for i in range(3)
                ]
            if verbose:
                print("Temperature:", t)
                print("Temperature color:", color)
            c.set_channel(LedChannel.FAN, LedMode.STATIC, b, color[0], color[1], color[2])
            
        if showCPUFreq:
            freqs = psutil.cpu_freq(percpu=True)
            max_freq = max(i.current for i in freqs)
            smoothed_freq = fSmoothing * smoothed_freq + (1 - fSmoothing) * max_freq
            interp_f = max(0, min(1, (smoothed_freq-2200)/(3700-2200)))            
            color = [
                int(interp_f * fHighColor[i] + (1 - interp_f)*fLowColor[i])
                for i in range(3)
                ]
            if verbose:
                print("Current Freq: ", max_freq)
                print("Smoothed Freq:", smoothed_freq)
                print("Frequency Color:", color)
            c.set_channel(LedChannel.LOGO, 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()

        time.sleep(interval)


if __name__ == '__main__':
    monitor()
