#!/usr/bin/python3

import time
import argparse
# import awinlib
# import awinlib.KasaSmartPowerStrip
import awinlib.KasaSmartPowerStrip as Kasa
from rich import print as pp
from prometheus_client import start_http_server, Gauge


ENABLE_EXPORTER = True
DEFAULT_PORT = 8088
DEFAULT_POLLING_INTERVAL = 5
DEFAULT_CONFIGS = {
    "port": DEFAULT_PORT,
    "polling_interval": DEFAULT_POLLING_INTERVAL
}


class KasaPowerStrip:
    def __init__(self, areaName, ip):
        self.areaName = areaName
        self.ip = ip
        try:
            self.powerStrip = Kasa.SmartPowerStrip(ip)
        except:
            self.powerStrip = None

    def collectData(self):
        dataDict = {}
        devSysInfo = self.powerStrip.get_system_info()
        for c in devSysInfo['system']['get_sysinfo']['children']:
            try:
                cid = int(c['id']) + 1
            except:
                continue

            # aliasName = c['alias'].encode('latin1').decode('UTF-8')
            gaugeName_Current_ma = f"{self.areaName}_{cid}_current_ma"
            gaugeName_Current_ma_desc = f"Current(milliampere) of {self.areaName}:{cid}"
            gaugeName_Voltage_mv = f"{self.areaName}_{cid}_voltage_mv"
            gaugeName_Voltage_mv_desc = f"Voltage(millivolt) of {self.areaName}:{cid}"
            gaugeName_Power_mw = f"{self.areaName}_{cid}_power_mw"
            gaugeName_Power_mw_desc = f"Power(milliwatt) of {self.areaName}:{cid}"
            gaugeName_TotalWh = f"{self.areaName}_{cid}_total_wh"
            gaugeName_TotalWh_desc = f"Total watt-hour of {self.areaName}:{cid}"

            ## Get Data
            realtimeInfo = self.powerStrip.get_realtime_energy_info(plug_num=cid)
            realTime_current_ma = realtimeInfo.get("current_ma", 0)
            realTime_voltage_mv = realtimeInfo.get("voltage_mv", 0)
            realTime_power_mw = realtimeInfo.get("power_mw", 0)
            realTime_total_wh = realtimeInfo.get("total_wh", 0)

            pp(f"Set {gaugeName_Current_ma}: {realTime_current_ma}")
            pp(f"Set {gaugeName_Voltage_mv}: {realTime_voltage_mv}")
            pp(f"Set {gaugeName_Power_mw}: {realTime_power_mw}")
            pp(f"Set {gaugeName_TotalWh}: {realTime_total_wh}")

            dataDict[cid] = {
                gaugeName_Current_ma: {
                    "description": gaugeName_Current_ma_desc,
                    "value": realTime_current_ma,
                },
                gaugeName_Voltage_mv: {
                    "description": gaugeName_Voltage_mv_desc,
                    "value": realTime_voltage_mv,
                },
                gaugeName_Power_mw: {
                    "description": gaugeName_Power_mw_desc,
                    "value": realTime_power_mw,
                },
                gaugeName_TotalWh: {
                    "description": gaugeName_TotalWh_desc,
                    "value": realTime_total_wh,
                }
            }
        return dataDict


def main(configs):
    time.sleep(10)
    gaugeDict = {}
    pollingInterval = configs.get("polling_interval", DEFAULT_POLLING_INTERVAL)
    start_http_server(configs.get("port", DEFAULT_PORT)) if ENABLE_EXPORTER else None

    powerStripList = [
        KasaPowerStrip("livingroom", "192.168.1.203"),
        KasaPowerStrip("attic", "192.168.1.203")
    ]

    while True:
        try:
            # collectData(gaugeDict)
            for powerStrip in powerStripList:
                powerStripData = powerStrip.collectData()
                for keyId in sorted(powerStripData.keys()):
                    for gaugeName in powerStripData[keyId].keys():
                        gaugeDesc = powerStripData[keyId][gaugeName]["description"]
                        gaugeValue = powerStripData[keyId][gaugeName]["value"]

                        if gaugeName not in gaugeDict:
                            gaugeDict[gaugeName] = Gauge(gaugeName, gaugeDesc)

                        if ENABLE_EXPORTER:
                            gaugeDict[gaugeName].set(gaugeValue)

                        pp(f"Set {gaugeName}: {gaugeValue}")

            time.sleep(pollingInterval)
        except KeyboardInterrupt:
            print("\nUser stopped process.")
            break


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-p", "--port", type=int, default=DEFAULT_PORT, help="The port of exportor")
    parser.add_argument("-i", "--polling_interval", default=DEFAULT_POLLING_INTERVAL, help="Polling interval for collect HDD temperature")
    args = parser.parse_args()
    DEFAULT_CONFIGS["port"] = args.port
    DEFAULT_CONFIGS["polling_interval"] = args.polling_interval
    main(DEFAULT_CONFIGS)