oin/oin_thermostat/mqtt.py

import json
import logging
import sys
from collections.abc import Callable
from typing import Any

import paho.mqtt.client as mqtt

from .screen import Screen
from .select import Selector

logger = logging.getLogger(__name__)


class HAClient:
    def __init__(
        self,
        entity: str,
        secondary_entities: list[str] = [],
        mqtt_config: dict[str, str] = dict(),
    ) -> None:
        self.entity = entity
        self.secondary_entities = secondary_entities
        self.config = mqtt_config

        self.state_topic = "oin/state"
        self.availability_topic = "oin/availability"

        self.client = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2)
        username = self.config.get("username", None)
        logger.debug(f"Setting up MQTT with user <{username}>")
        self.client.username_pw_set(
            username=username,
            password=self.config.get("password", None),
        )

        self.screen = Screen()
        self.selector = Selector(self.send_data)

    @property
    def ha_options(self) -> dict[str, str | dict[str, str]]:
        return {
            "dev": {
                "ids": "oin",
                "name": "Oin",
            },
            "o": {
                "name": "Oin",
            },
            "availability_topic": self.availability_topic,
            "state_topic": self.state_topic,
            "cmps": self.selector.ha_options,
        }

    def connect(self) -> None:
        self.client.will_set(self.availability_topic, "offline", retain=True)

        host = self.config.get("host")
        port = self.config.get("port", 1883)
        logger.debug(f"Connecting to <{host}> on port <{port}>.")
        self.client.connect(host, port)

        self.subscribe(entity_topic(self.entity), self.state_update)
        self.subscribe(
            [entity_topic(entity) for entity in self.secondary_entities],
            self.secondary_state_update,
        )

        self.publish("homeassistant/device/oin/config", self.ha_options, retain=True)
        self.client.publish(self.availability_topic, "online", retain=True)

    def publish(self, topic: str, data: Any, **kwargs) -> mqtt.MQTTMessageInfo:
        logger.debug(f"Sending message on topic <{topic}>: {json.dumps(data)}")
        return self.client.publish(topic, json.dumps(data), **kwargs)

    def subscribe(self, topic: str | list[str], callback: Callable) -> None:
        logger.debug(f"Subscribing to <{topic}>.")

        match topic:
            case str():
                self.client.message_callback_add(topic, callback)
                code, _ = self.client.subscribe(topic)
            case list():
                for top in topic:
                    self.client.message_callback_add(top, callback)
                code, _ = self.client.subscribe([(top, 0) for top in topic])

        if code != 0:
            logger.error(f"Failed subscribing to topic <{topic}> with code <{code}>.")
            sys.exit(1)

    def loop(self) -> mqtt.MQTTErrorCode:
        logger.info("Starting MQTT client loop.")
        code = self.client.loop_forever(retry_first_connection=True)

        if code != 0:
            logger.error("MQTT client loop failed with code <{code}>.")

    def state_update(
        self, client: mqtt.Client, userdata: Any, message: mqtt.MQTTMessage
    ) -> None:
        logger.debug(f"Message received on topic <{message.topic}>: {message.payload}.")

        subtopic = message.topic.rsplit("/", maxsplit=1)[1]

        match subtopic:
            case "current_temperature":
                self.screen.value = parse(message)
            case "temperature":
                if (value := parse(message)) != self.selector.temperature:
                    self.screen.tmp_value = value
                    self.selector.temperature = value
            case "hvac_action":
                self.screen.mode = parse(message)
            case "preset_modes":
                if (value := parse(message)) != self.selector.preset_modes:
                    self.selector.preset_modes = value
            case "preset_mode":
                if (value := parse(message)) != self.selector.mode:
                    self.selector.mode = value
            case "state":
                match message.payload.decode():
                    case "heat":
                        self.selector.switch = True
                    case "off":
                        self.selector.switch = False

    def secondary_state_update(
        self, client: mqtt.Client, userdata: Any, message: mqtt.MQTTMessage
    ) -> None:
        logger.debug(f"Message received on topic <{message.topic}>: {message.payload}.")

        _, grp, ent, subtopic = message.topic.split("/")
        idx = self.secondary_entities.index(f"{grp}.{ent}")

        if subtopic == "state":
            self.screen.secondary |= {idx: message.payload.decode()}

    def send_data(self, data: Any) -> mqtt.MQTTMessageInfo:
        return self.publish(self.state_topic, data)


def parse(message: mqtt.MQTTMessage) -> Any:
    return json.loads(message.payload.decode())


def entity_topic(entity: str, subtopic: str = "#") -> str:
    topic = entity.replace(".", "/")
    return f"homeassistant/{topic}/{subtopic}"