mqttdevicemanager/mqtthandler/handler.py

import aiomqtt
import asyncio
import inspect
import paho
import signal
import json
from dataclasses import dataclass

from .command import (
    Command,
    CommandResponse,
    CommandArgumentError,
    CommandExecutionError,
)


@dataclass
class MQTTConfig:
    host: str
    port: int = 1883
    username: str | None = None
    password: str | None = None
    keepalive: int = 60


class MQTTHandler:
    STATUS = "status"

    def __init__(
        self,
        mqtt_config: MQTTConfig,
        handler_id: str,
    ):
        self.handler_id = handler_id
        self.mqtt_config = mqtt_config

        self.topic_base = f"device/{handler_id}"
        self.command_topic = f"{self.topic_base}/command"
        self.property_topic = f"{self.topic_base}/property"

        self._shutdown_event = asyncio.Event()
        will = aiomqtt.Will(
            topic=f"{self.property_topic}/{MQTTHandler.STATUS}", payload="OFFLINE", qos=1, retain=True
        )

        self.mqtt_client = aiomqtt.Client(
            self.mqtt_config.host,
            port=self.mqtt_config.port,
            identifier=handler_id,
            protocol=paho.mqtt.client.MQTTv5,
            will=will,
        )

        self.commands = self.get_available_commands()

    def get_available_commands(self):
        commands = {}
        for base in self.__class__.__mro__:
            for name, attr in vars(base).items():
                if isinstance(attr, Command):
                    print(f"Registering method {type(self).__name__}.{name} as command '{attr.name}'")
                    commands[attr.name] = attr

        return commands

    async def register_commands(self):
        for name, command in self.commands.items():
            for k, v in {
                "schema": json.dumps(command.schema),
                "description": command.description,
                **command.additional_properties,
            }.items():
                await self.mqtt_client.publish(
                    f"{self.command_topic}/{command.name}/${k}",
                    str(v),
                    qos=1,
                    retain=True,
                )

    async def register_property(self, property: str, description: str | None = None, schema: dict | None = None):
        data = {
            "schema": json.dumps(schema),
            "description": description,
        }
        for k, v in {k:v for k,v in data.items() if v is not None}.items():
            await self.mqtt_client.publish(
                f"{self.property_topic}/{property}/${k}",
                str(v),
                qos=1,
                retain=True,
            )

    async def set_property(self, property: str, value, qos=0, retain=False):
        await self.mqtt_client.publish(
            f"{self.property_topic}/{property}",
            str(value),
            qos=qos,
            retain=retain,
        )

    async def execute_command(
        self,
        command_name: str,
        payload: str,
        properties: paho.mqtt.properties.Properties = None,
    ):
        async def respond(success: bool, message: str = None):
            if (
                properties is not None
                and hasattr(properties, "ResponseTopic")
                and properties.ResponseTopic is not None
            ):
                correlation = (
                    properties.CorrelationData.decode("utf-8")
                    if hasattr(properties, "CorrelationData")
                    else None
                )
                await self.mqtt_client.publish(
                    properties.ResponseTopic,
                    str(CommandResponse(success, str(message), correlation)),
                    qos=1,
                    retain=False,
                )

        try:
            command = self.commands[command_name]
            result = await command(self, payload)
            await respond(True, result)
        except (CommandArgumentError, CommandExecutionError) as e:
            await respond(False, f"{e}")
        except Exception as e:
            print(f"Failed to execute command {command_name} with unknown cause: ", e)
            await respond(False, "Unexpected error")

    async def mqtt_command_writer_task(self):
        async for message in self.mqtt_client.messages:
            topic = str(message.topic)
            payload = message.payload.decode("utf-8")

            if topic.startswith(self.command_topic):
                command_name = topic.removeprefix(f"{self.command_topic}/")
                await self.execute_command(command_name, payload, message.properties)

    async def shutdown_watcher(self):
        await self._shutdown_event.wait()
        await self.set_property(MQTTHandler.STATUS, "OFFLINE", qos=1, retain=True)

    def stop(self):
        self._shutdown_event.set()
        signal.signal(signal.SIGINT, signal.SIG_DFL)

    async def run(self):
        INTERVAL = 5
        while True:
            try:
                async with self.mqtt_client as client:
                    await client.subscribe(f"{self.command_topic}/+")
                    await self.register_commands()

                    # announce that we are online
                    await self.set_property(MQTTHandler.STATUS, "ONLINE", qos=1, retain=True)
                    await self.register_property("status", "Indicates the status of the device.", {
                        "type": "string",
                        "enum": ["ONLINE", "OFFLINE"]
                    })

                    tasks = [self.mqtt_command_writer_task(), self.shutdown_watcher()]

                    # Inspect instance methods
                    for attr_name in dir(self):
                        attr = getattr(self, attr_name)
                        if callable(attr) and getattr(attr, "_is_task", False):
                            if not inspect.iscoroutinefunction(attr):
                                raise TypeError(
                                    f"@task can only decorate async methods: {attr_name}"
                                )
                            tasks.append(attr())  # call it, returns coroutine

                    await asyncio.gather(*tasks)

            except aiomqtt.MqttError as e:
                print(
                    f"[{self.handler_id}] MQTT connection error: {e}. Reconnecting in {INTERVAL}s..."
                )
                await asyncio.sleep(INTERVAL)


def task(func):
    """Decorator to mark async methods for automatic gathering."""
    func._is_task = True
    return func