import logging
from sys import argv
from time import sleep

import paho.mqtt.client as mqtt
from pyModbusTCP.client import ModbusClient
from pyModbusTCP.utils import decode_ieee, word_list_to_long

logging.basicConfig(level=logging.DEBUG)
LOG = logging.getLogger(__name__)

class SDM630_MQTT:
    def __init__(
        self,
        modbus_host,
        mqtt_host="localhost",
        modbus_port=4196,
        mqtt_port=1883,
        modbus_unit=1,
        mqtt_prefix="sdm630",
    ):
        self.modbus = ModbusClient(
            host=modbus_host,
            port=modbus_port,
            unit_id=modbus_unit,
        )
        self.modbus.timeout = 10

        self.mqtt = mqtt.Client()
        self.mqtt_host = mqtt_host
        self.mqtt_port = mqtt_port
        self.mqtt_prefix = mqtt_prefix

        self.mqtt.on_connect = self._on_mqtt_connect
        self.mqtt.on_disconnect = self._on_mqtt_disconnect

    def _publish(self, topic, msg):
        LOG.debug(f"_publish({topic!r}, {msg!r})")
        self.mqtt.publish(f"{self.mqtt_prefix.rstrip('/')}/{topic}", msg)

    def _on_mqtt_connect(self, client, userdata, flags, rc):
        LOG.info(f"Connected to mqtt server {self.mqtt_host}:{self.mqtt_port}")

    def _on_mqtt_disconnect(self, client, userdata, rc):
        LOG.info(f"Disconnected from mqtt server {self.mqtt_host}:{self.mqtt_port}")

    def start(self):
        LOG.debug("_start()")
        self.mqtt.connect(self.mqtt_host, self.mqtt_port, 10)
        self.mqtt.loop_start()

    def stop(self):
        self.mqtt.loop_stop()
        self.mqtt.disconnect()

    def request_and_publish(self):
        self.modbus.open()
        for idx, (readable, address) in enumerate(
            {
                "active-power/L1": 0x000C,
                "active-power/L2": 0x000E,
                "active-power/L3": 0x0010,
                "active-power/total": 0x0034,
                "apparent-power/L1": 0x0012,
                "apparent-power/L2": 0x0014,
                "apparent-power/L3": 0x0016,
                "apparent-power/total": 0x0038,
                "current-demand/L1": 0x0102,
                "current-demand/L2": 0x0104,
                "current-demand/L3": 0x0106,
                "current-demand/N": 0x0068,
                "current-demand/max_L1": 0x0108,
                "current-demand/max_L2": 0x010A,
                "current-demand/max_L3": 0x010C,
                "current-demand/max_N": 0x006A,
                "current/L1": 0x0006,
                "current/L2": 0x0008,
                "current/L3": 0x000A,
                "current/average": 0x002E,
                "current/neutral": 0x00E0,
                "current/total": 0x0030,
                "export_kVArh": 0x004C,
                "export_kWh": 0x004A,
                "frequency": 0x0046,
                "import_kVArh": 0x004E,
                "import_kWh": 0x0048,
                "phase-angle/L1": 0x0024,
                "phase-angle/L2": 0x0026,
                "phase-angle/L3": 0x0028,
                "phase-angle/total": 0x003E,
                "power-demand/max": 0x0056,
                "power-demand/total": 0x0054,
                "power-factor/L1": 0x001E,
                "power-factor/L2": 0x0020,
                "power-factor/L3": 0x0022,
                "power-factor/total": 0x003E,
                "reactive-power/L1": 0x0018,
                "reactive-power/L2": 0x001A,
                "reactive-power/L3": 0x001E,
                "reactive-power/total": 0x003C,
                "voltage/L1_L2": 0x00C8,
                "voltage/L1_N": 0x0000,
                "voltage/L2_L3": 0x00CA,
                "voltage/L2_N": 0x0002,
                "voltage/L3_L1": 0x00CC,
                "voltage/L3_N": 0x0004,
                "voltage/average": 0x002A,
            }.items()
        ):
            read = self.modbus.read_input_registers(address, 2)
            longs = word_list_to_long(read)
            if len(longs) == 1:
                self._publish(readable, decode_ieee(longs[0]))
            else:
                self._publish(readable, sum([decode_ieee(x) for x in longs]))
            if idx % 20:
                self.modbus.close()
                self.modbus.open()
        self.modbus.close()


if __name__ == "__main__":
    client = SDM630_MQTT(
        "192.168.1.254",
        "127.0.0.1",
    )
    client.start()
    try:
        while True:
            client.request_and_publish()
            sleep(2)
    except Exception as e:
        LOG.exception('oops')
        client.stop()