srne-mqtt/srnemqtt/consumers/mqtt.py

# -*- coding: utf-8 -*-
import json
from logging import getLogger
from time import sleep
from typing import Any, Dict, List, Optional, TypeAlias
from uuid import uuid4

import paho.mqtt.client as mqtt

from ..solar_types import DataName
from . import BaseConsumer

logger = getLogger(__name__)

MAP_VALUES: Dict[DataName, Dict[str, Any]] = {
    # DataName.BATTERY_VOLTAGE_MIN: {},
    # DataName.BATTERY_VOLTAGE_MAX: {},
    # DataName.CHARGE_MAX_CURRENT: {},
    # DataName.DISCHARGE_MAX_CURRENT: {},
    # DataName.CHARGE_MAX_POWER: {},
    # DataName.DISCHARGE_MAX_POWER: {},
    # DataName.CHARGE_AMP_HOUR: {},
    # DataName.DISCHARGE_AMP_HOUR: {},
    DataName.PRODUCTION_ENERGY: {
        "unit": "Wh",
        "type": "energy",
        "state_class": "total_increasing",
    },
    DataName.CONSUMPTION_ENERGY: {
        "unit": "Wh",
        "type": "energy",
        "state_class": "total_increasing",
    },
    # DataName.RUN_DAYS: {},
    # DataName.DISCHARGE_COUNT: {},
    # DataName.FULL_CHARGE_COUNT: {},
    # DataName.TOTAL_CHARGE_AMP_HOURS: {},
    # DataName.TOTAL_DISCHARGE_AMP_HOURS: {},
    DataName.TOTAL_PRODUCTION_ENERGY: {
        "unit": "Wh",
        "type": "energy",
        "state_class": "total_increasing",
        "expiry": 180,
    },
    DataName.TOTAL_CONSUMPTION_ENERGY: {
        "unit": "Wh",
        "type": "energy",
        "state_class": "total_increasing",
        "expiry": 180,
    },
    #
    DataName.BATTERY_CHARGE: {
        "unit": "%",
        "type": "battery",
        "state_class": "measurement",
    },
    DataName.BATTERY_VOLTAGE: {
        "unit": "V",
        "type": "voltage",
        "state_class": "measurement",
    },
    DataName.BATTERY_CURRENT: {
        "unit": "A",
        "type": "current",
        "state_class": "measurement",
    },
    DataName.INTERNAL_TEMPERATURE: {
        "unit": "°C",
        "type": "temperature",
        "state_class": "measurement",
    },
    DataName.BATTERY_TEMPERATURE: {
        "unit": "°C",
        "type": "temperature",
        "state_class": "measurement",
    },
    DataName.LOAD_VOLTAGE: {
        "unit": "V",
        "type": "voltage",
        "state_class": "measurement",
    },
    DataName.LOAD_CURRENT: {
        "unit": "A",
        "type": "current",
        "state_class": "measurement",
    },
    DataName.LOAD_POWER: {"unit": "W", "type": "power", "state_class": "measurement"},
    DataName.LOAD_ENABLED: {
        "type": "outlet",
        "platform": "switch",
    },
    DataName.PANEL_VOLTAGE: {
        "unit": "V",
        "type": "voltage",
        "state_class": "measurement",
    },
    DataName.PANEL_CURRENT: {
        "unit": "A",
        "type": "current",
        "state_class": "measurement",
    },
    DataName.PANEL_POWER: {"unit": "W", "type": "power", "state_class": "measurement"},
    # DataName.LOAD_ENABLED: {},
    DataName.CALCULATED_BATTERY_POWER: {
        "unit": "W",
        "type": "power",
        "state_class": "measurement",
    },
    DataName.CALCULATED_PANEL_POWER: {
        "unit": "W",
        "type": "power",
        "state_class": "measurement",
    },
    DataName.CALCULATED_LOAD_POWER: {
        "unit": "W",
        "type": "power",
        "state_class": "measurement",
    },
}


PayloadType: TypeAlias = str | bytes | bytearray | int | float | None


class MqttConsumer(BaseConsumer):
    initialized: List[str]

    _client: mqtt.Client | None = None

    def __init__(self, settings: Dict[str, Any]) -> None:
        self.initialized = []

        super().__init__(settings)

    @property
    def client(self) -> mqtt.Client:
        if self._client is not None:
            return self._client

        self._client = mqtt.Client(
            client_id=self.settings["client"]["id"], userdata=self
        )
        self._client.on_connect = self.on_connect
        self._client.on_message = self.on_message
        self._client.on_disconnect = self.on_disconnect
        self._client.on_connect_fail = self.on_connect_fail
        # Will must be set before connecting!!
        self._client.will_set(
            f"{self.topic_prefix}/available", payload="offline", retain=True
        )
        while True:
            try:
                self._client.connect(
                    self.settings["client"]["host"],
                    self.settings["client"]["port"],
                    self.settings["client"]["keepalive"],
                )
                break
            except OSError as err:
                # Network is unreachable
                if err.errno == 101:
                    pass
                # Temporary failure in name resolution
                elif err.errno == -3:
                    pass
                else:
                    logger.exception("Unknown error connecting to mqtt server")
                    raise
                logger.warning(
                    "Temporary failure connecting to mqtt server", exc_info=True
                )
                sleep(0.1)
        return self._client

    def config(self, settings: Dict[str, Any]):
        super().config(settings)
        settings.setdefault("client", {})
        settings["client"].setdefault("id", None)
        settings["client"].setdefault("host", "")
        settings["client"].setdefault("port", 1883)
        settings["client"].setdefault("keepalive", 60)

        if not settings.get("device_id"):
            settings["device_id"] = str(uuid4())

        settings.setdefault("prefix", "solarmppt")

        settings.setdefault("discovery_prefix", "homeassistant")

    _controller_id: str | None = None

    @property
    def controller_id(self) -> str:
        assert self.controller is not None
        # Controller serial is fetched from device, cache it.
        if self._controller_id is None:
            self._controller_id = self.controller.serial
        return f"{self.controller.manufacturer_id}_{self._controller_id}"

    @property
    def topic_prefix(self):
        return f"{self.settings['prefix']}/{self.controller_id}"

    def get_ha_config(
        self,
        id: str,
        name: str,
        unit: Optional[str] = None,
        type: Optional[str] = None,
        expiry: int = 90,
        state_class: Optional[str] = None,
        platform: str = "sensor",
    ):
        assert state_class in [None, "measurement", "total", "total_increasing"]
        assert self.controller is not None

        res = {
            "~": f"{self.topic_prefix}",
            "unique_id": f"{self.controller_id}_{id}",
            "object_id": f"{self.controller_id}_{id}",  # Used for entity id
            "availability_topic": "~/available",
            "state_topic": f"~/{id}",
            "name": name,
            "device": {
                "identifiers": [
                    self.controller_id,
                ],
                "manufacturer": self.controller.manufacturer,
                "model": self.controller.model,
                "sw_version": self.controller.version,
                "via_device": self.settings["device_id"],
                "suggested_area": "Solar panel",
                "name": self.controller.name,
            },
            "force_update": True,
            "expire_after": expiry,
        }

        if unit:
            res["unit_of_meas"] = unit
        if type:
            res["dev_cla"] = type
        if state_class:
            res["state_class"] = state_class
        if platform == "switch":
            res["command_topic"] = f"{res['state_topic']}/set"
            res["payload_on"] = True
            res["payload_off"] = False
        return res

    # The callback for when the client receives a CONNACK response from the server.
    @staticmethod
    def on_connect(client: mqtt.Client, userdata: "MqttConsumer", flags, rc):
        logger.info("MQTT connected with result code %s", rc)

        # Subscribing in on_connect() means that if we lose the connection and
        # reconnect then subscriptions will be renewed.
        # client.subscribe("$SYS/#")
        client.publish(
            f"{userdata.topic_prefix}/available", payload="online", retain=True
        )

        load_set_topic = f"{userdata.topic_prefix}/load_enabled/set"
        client.message_callback_add(load_set_topic, userdata.on_load_switch)
        client.subscribe(load_set_topic)

    @staticmethod
    def on_load_switch(
        client: mqtt.Client, userdata: "MqttConsumer", message: mqtt.MQTTMessage
    ):
        assert userdata.controller is not None
        logger.debug(message.payload)
        payload = message.payload.decode().upper() in ("ON", "TRUE", "ENABLE", "YES")

        res = userdata.controller.load_enabled = payload
        client.publish(
            f"{userdata.topic_prefix}/load_enabled", payload=res, retain=True
        )

    @staticmethod
    def on_connect_fail(client: mqtt.Client, userdata: "MqttConsumer"):
        logger.warning("on_connect_fail")

    # The callback for when a PUBLISH message is received from the server.
    @staticmethod
    def on_message(client, userdata, msg):
        logger.info(msg.topic + " " + str(msg.payload))

    @staticmethod
    def on_disconnect(client: mqtt.Client, userdata: "MqttConsumer", rc, prop=None):
        logger.warning("on_disconnect %s", rc)

    def poll(self):
        res = self.client.loop(timeout=0.1, max_packets=5)

        if res != mqtt.MQTT_ERR_SUCCESS:
            logger.warning("loop returned non-success: %s", res)
            try:
                sleep(1)
                res = self.client.reconnect()
                if res != mqtt.MQTT_ERR_SUCCESS:
                    logger.error("Reconnect failed: %s", res)
            except (OSError, mqtt.WebsocketConnectionError) as err:
                logger.error("Reconnect failed: %s", err)

        return super().poll()

    def write(self, data: Dict[str, PayloadType]):
        self.client.publish(f"{self.topic_prefix}/raw", payload=json.dumps(data))

        for dataname, data_value in data.items():
            if dataname in MAP_VALUES:
                if dataname not in self.initialized:
                    km = MAP_VALUES[DataName(dataname)]
                    pretty_name = dataname.replace("_", " ").capitalize()
                    disc_prefix = self.settings["discovery_prefix"]
                    platform = km.get("platform", "sensor")

                    self.client.publish(
                        f"{disc_prefix}/{platform}/{self.controller_id}/{dataname}/config",
                        payload=json.dumps(
                            self.get_ha_config(dataname, pretty_name, **km)
                        ),
                        retain=True,
                    )
                    self.initialized.append(dataname)

                self.client.publish(
                    f"{self.topic_prefix}/{dataname}", data_value, retain=True
                )

    def exit(self):
        self.client.publish(
            f"{self.topic_prefix}/available", payload="offline", retain=True
        )

        while self.client.want_write():
            self.client.loop_write(10)

        self.client.disconnect()
        return super().exit()


# Client(client_id="", clean_session=True, userdata=None,
# protocol=MQTTv311, transport="tcp")

# connect_srv(domain, keepalive=60, bind_address="")
# Connect to a broker using an SRV DNS lookup to obtain the broker address.
# Takes the following arguments:
# domain
#    the DNS domain to search for SRV records.
# If None, try to determine the local domain name.

# client.will_set(topic, payload=None, qos=0, retain=False)

# Blocking call that processes network traffic, dispatches callbacks and
# handles reconnecting.
# Other loop*() functions are available that give a threaded interface and a
# manual interface.

# client.loop_forever()