home-automation/app/integrations/mqtt.py

"""MQTT client integration (paho-mqtt 2.x).

Provides :class:`MqttManager`: a long-lived MQTT client that runs paho's
loop in a background thread.  Designed to be started in FastAPI's lifespan
and shared across the application as a module-level singleton (``mqtt_manager``).

Key design decisions
--------------------
- **paho 2.x API**: ``Client`` requires ``CallbackAPIVersion.VERSION2`` as the
  first positional argument.  VERSION2 on_connect callback receives
  ``(client, userdata, connect_flags, reason_code, properties)`` — not the
  older RC int.
- **Background thread**: ``loop_start()`` spawns a daemon thread; ``loop_stop()``
  joins it cleanly on shutdown.
- **Never crash the process**: all paho operations are wrapped in try/except;
  connection failure is logged but does not propagate to the caller.
- **Password safety**: the MQTT password is never passed to ``logger`` calls.
- **Reconnect**: ``reconnect(settings)`` tears down the old client and
  establishes a fresh connection with the new settings.  Callers (e.g. PUT
  /api/config) call this after saving MQTT-related config values.
- **Subscribe**: ``subscribe(topic, handler)`` registers a topic → handler
  mapping.  Subscriptions are re-established automatically on reconnect.
  The ``on_message`` callback dispatches incoming payloads to the registered
  handler; handler exceptions are caught and logged so they never crash the
  paho loop thread or the network connection.
"""

from __future__ import annotations

import logging
import threading
from collections.abc import Callable
from dataclasses import dataclass
from typing import TYPE_CHECKING

import paho.mqtt.client as mqtt

if TYPE_CHECKING:
    from app.config import Settings

logger = logging.getLogger(__name__)

# MQTT settings keys that, when changed, should trigger a reconnect.
MQTT_SETTINGS_KEYS = {
    "mqtt_enabled",
    "mqtt_broker_host",
    "mqtt_broker_port",
    "mqtt_username",
    "mqtt_password",
    "mqtt_tls_enabled",
    "mqtt_client_id",
}


@dataclass
class _SourceClientState:
    """One installed source generation and its in-flight callback count."""

    client: mqtt.Client
    generation: int
    in_flight: int = 0


def _is_configured(settings: Settings) -> bool:
    """Return True if MQTT is enabled *and* the broker host is set."""
    return bool(settings.mqtt_enabled and settings.mqtt_broker_host)


def mqtt_source_client_id(base_client_id: str, source_id: int) -> str:
    """Return a stable, deployment-scoped identity for one DSMR source."""
    return f"{base_client_id}-dsmr-source-{source_id}"


def mqtt_test_client_id(base_client_id: str) -> str:
    """Return a transient test identity that cannot evict a long-lived client."""
    return f"{base_client_id}-test"


class MqttManager:
    """Long-lived MQTT client wrapper.

    Lifecycle
    ---------
    1. ``connect(settings)`` — start background loop + connect to broker.
    2. ``publish(...)`` — publish messages while connected.
    3. ``disconnect()`` — graceful teardown (called on app shutdown).
    4. ``reconnect(settings)`` — disconnect then re-connect with new settings
       (called after saving updated MQTT configuration).

    When MQTT is not enabled or not configured, all methods are no-ops.
    Connection failures are caught and logged; they do not raise.
    """

    def __init__(self) -> None:
        self._client: mqtt.Client | None = None
        self._lock = threading.RLock()
        # Source replacement/removal is serialized independently from callback
        # bookkeeping.  In particular, paho's loop_stop() joins its network
        # thread, whose callback completion also needs ``_lock``.
        self._source_lifecycle_lock = threading.Lock()
        self._source_idle = threading.Condition(self._lock)
        self._connected = False
        # topic → handler registry; persists across reconnects so subscriptions
        # are automatically re-established when the client reconnects.
        self._subscriptions: dict[str, Callable[[bytes], None]] = {}
        # DSMR sources are independent connections: their credentials and TLS
        # configuration belong to MeterSource.config, not app_config.
        self._source_clients: dict[int, mqtt.Client] = {}
        self._source_subscriptions: dict[int, dict[str, Callable[[bytes], None]]] = {}
        self._source_connected: set[int] = set()
        # Each replacement gets a distinct identity.  A paho callback can run
        # after its client was stopped, so source id alone is not sufficient.
        self._source_generations: dict[int, int] = {}
        self._source_states: dict[int, _SourceClientState] = {}
        self._next_source_generation = 0

    # ------------------------------------------------------------------
    # Public properties
    # ------------------------------------------------------------------

    @property
    def is_connected(self) -> bool:
        """True if the underlying paho client is currently connected."""
        return self._connected

    # ------------------------------------------------------------------
    # Public interface
    # ------------------------------------------------------------------

    def is_configured(self, settings: Settings) -> bool:
        """Return True if MQTT is enabled and broker host is configured."""
        return _is_configured(settings)

    def connect(self, settings: Settings) -> None:
        """Connect to the MQTT broker and start the background loop thread.

        No-op if MQTT is not enabled or broker host is not set.
        Connection errors are logged but do not raise.
        """
        if not _is_configured(settings):
            logger.debug("MQTT not configured or not enabled — skipping connect.")
            return

        with self._lock:
            self._start_client(settings)

    def disconnect(self) -> None:
        """Disconnect from the broker and stop the background loop thread.

        No-op if no client is active.
        """
        with self._lock:
            self._stop_client()
        with self._source_lifecycle_lock:
            with self._lock:
                source_ids = list(self._source_clients)
            for source_id in source_ids:
                self._stop_source_client(source_id)

    def reconnect(self, settings: Settings) -> None:
        """Disconnect the current client (if any) and reconnect with *settings*.

        Call this after saving updated MQTT configuration values.
        No-op if MQTT is not enabled or broker host is not set.
        """
        with self._lock:
            self._stop_client()
        self.connect(settings)

    def publish(
        self,
        topic: str,
        payload: str | bytes | None,
        *,
        retain: bool = False,
        qos: int = 0,
    ) -> bool:
        """Publish a message to *topic*.

        Return whether paho accepted the message for publication.  If the
        client is not connected, or paho raises/rejects it, return ``False``;
        errors still do not escape this best-effort boundary.
        """
        with self._lock:
            client = self._client

        if client is None or not self._connected:
            logger.debug("MQTT publish skipped — not connected (topic=%s).", topic)
            return False

        try:
            result = client.publish(topic, payload=payload, qos=qos, retain=retain)
            # paho returns MQTTMessageInfo with an integer rc.  Keep fakes and
            # alternate clients compatible by treating an absent/non-int rc as
            # accepted after the call itself succeeded.
            rc = getattr(result, "rc", None)
            if isinstance(rc, int) and rc != mqtt.MQTT_ERR_SUCCESS:
                logger.warning("MQTT publish rejected (topic=%s, rc=%s).", topic, rc)
                return False
            return True
        except Exception:
            logger.exception("MQTT publish error (topic=%s).", topic)
            return False

    def subscribe(self, topic: str, handler: Callable[[bytes], None]) -> None:
        """Register *handler* to be called when a message arrives on *topic*.

        If the client is already connected, the subscription is sent to the
        broker immediately.  Otherwise it is queued and will be established
        the next time ``_on_connect`` fires (including after a reconnect).

        Handler exceptions are swallowed by the ``on_message`` dispatcher so
        that a buggy handler can never crash the paho loop thread.
        """
        self._subscriptions[topic] = handler
        # Grab the client reference outside the lock to avoid holding the lock
        # while calling back into paho (which may itself acquire internal locks).
        with self._lock:
            client = self._client
            connected = self._connected

        if client is not None and connected:
            try:
                client.subscribe(topic)
                logger.debug("MQTT subscribed to topic=%s (immediate).", topic)
            except Exception:
                logger.exception("MQTT subscribe error (topic=%s).", topic)

    def unsubscribe(self, topic: str) -> None:
        """Remove the handler for *topic* and unsubscribe from the broker.

        Idempotent: unknown topics are ignored.  Used to apply config changes
        without a restart (e.g. when DSMR ingest is turned off or its topic
        changes).  If the client is connected, the broker unsubscribe is sent
        immediately; either way the handler is removed from the registry so it
        will not be re-subscribed on the next ``_on_connect``.
        """
        self._subscriptions.pop(topic, None)
        with self._lock:
            client = self._client
            connected = self._connected

        if client is not None and connected:
            try:
                client.unsubscribe(topic)
                logger.debug("MQTT unsubscribed from topic=%s.", topic)
            except Exception:
                logger.exception("MQTT unsubscribe error (topic=%s).", topic)

    def replace_source(
        self,
        source_id: int,
        *,
        host: str,
        port: int,
        username: str,
        password: str,
        tls_enabled: bool,
        subscriptions: dict[str, Callable[[bytes], None]],
        base_client_id: str = "home-automation",
        state_handler: Callable[[str], None] | None = None,
    ) -> bool:
        """Replace one source-owned client and its handlers.

        This intentionally does not touch the legacy app-wide client or any
        other source client.  It is also safe for a source to be temporarily
        unconfigured: handlers are retained in the source registry but no
        connection is attempted until a host is supplied.
        """
        with self._source_lifecycle_lock:
            self._stop_source_client(source_id)
            if not host:
                self._report_source_state(state_handler, "error", source_id)
                return False
            self._next_source_generation += 1
            generation = self._next_source_generation
            captured_subscriptions = dict(subscriptions)
            client = mqtt.Client(
                callback_api_version=mqtt.CallbackAPIVersion.VERSION2,
                client_id=mqtt_source_client_id(base_client_id, source_id),
            )

            def _on_connect(
                connected_client: mqtt.Client,
                _userdata: object,
                _flags: mqtt.ConnectFlags,
                reason_code: mqtt.ReasonCode,
                _properties: mqtt.Properties | None,
            ) -> None:
                with self._lock:
                    if not self._is_current_source_client(source_id, generation, connected_client):
                        return
                    if reason_code.is_failure:
                        self._source_connected.discard(source_id)
                        logger.warning("DSMR MQTT connection refused for source_id=%s", source_id)
                        state = "error"
                    else:
                        self._source_connected.add(source_id)
                        state = "online"
                        for topic in captured_subscriptions:
                            try:
                                connected_client.subscribe(topic)
                            except Exception:
                                logger.exception("DSMR MQTT re-subscribe failed for source_id=%s", source_id)
                self._report_source_state(state_handler, state, source_id)

            def _on_disconnect(
                disconnected_client: mqtt.Client,
                _userdata: object,
                _flags: mqtt.DisconnectFlags,
                _reason_code: mqtt.ReasonCode,
                _properties: mqtt.Properties | None,
            ) -> None:
                with self._lock:
                    if not self._is_current_source_client(source_id, generation, disconnected_client):
                        return
                    self._source_connected.discard(source_id)
                self._report_source_state(state_handler, "error", source_id)

            def _on_message(
                message_client: mqtt.Client,
                _userdata: object,
                message: mqtt.MQTTMessage,
            ) -> None:
                with self._lock:
                    if not self._is_current_source_client(source_id, generation, message_client):
                        return
                    handler = captured_subscriptions.get(message.topic)
                    state = self._source_states.get(source_id)
                    if handler is None or state is None:
                        return
                    # This permit covers the entire handler call. Teardown first
                    # invalidates the state and then waits for all permits, so an
                    # old callback cannot run after teardown returns.
                    state.in_flight += 1
                try:
                    handler(message.payload)
                except Exception:
                    logger.exception("DSMR source handler raised (source_id=%s)", source_id)
                finally:
                    with self._lock:
                        state.in_flight -= 1
                        if state.in_flight == 0:
                            self._source_idle.notify_all()

            client.on_connect = _on_connect
            client.on_disconnect = _on_disconnect
            client.on_message = _on_message
            if tls_enabled:
                try:
                    client.tls_set()
                except Exception:
                    logger.exception("DSMR MQTT TLS setup failed for source_id=%s", source_id)
                    self._report_source_state(state_handler, "error", source_id)
                    return False
            if username:
                client.username_pw_set(username=username, password=password or None)
            # Register ownership before network processing begins.  A broker
            # may deliver CONNACK synchronously from connect(), or on the loop
            # thread before connect() returns.
            with self._lock:
                self._source_clients[source_id] = client
                self._source_subscriptions[source_id] = captured_subscriptions
                self._source_generations[source_id] = generation
                self._source_states[source_id] = _SourceClientState(client, generation)
            self._report_source_state(state_handler, "connecting", source_id)
            client.loop_start()
            try:
                client.connect(host=host, port=port, keepalive=60)
            except Exception:
                logger.exception("DSMR MQTT connect failed (source_id=%s, host=%s)", source_id, host)
                self._report_source_state(state_handler, "error", source_id)
                self._stop_source_client(source_id)
                return False
            return True

    def remove_source(self, source_id: int) -> None:
        """Drop one source client and its handlers, including queued callbacks."""
        with self._source_lifecycle_lock:
            self._stop_source_client(source_id)

    def source_is_active(self, source_id: int) -> bool:
        """Whether a source-owned client is currently installed for callbacks."""
        with self._lock:
            return source_id in self._source_clients

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _start_client(self, settings: Settings) -> None:
        """Build a fresh paho Client, configure it, and call loop_start + connect."""
        client = mqtt.Client(
            callback_api_version=mqtt.CallbackAPIVersion.VERSION2,
            client_id=settings.mqtt_client_id,
        )

        # Callbacks — VERSION2 on_connect signature:
        # (client, userdata, connect_flags, reason_code, properties)
        def _on_connect(
            _client: mqtt.Client,
            _userdata: object,
            _flags: mqtt.ConnectFlags,
            reason_code: mqtt.ReasonCode,
            _properties: mqtt.Properties | None,
        ) -> None:
            if reason_code.is_failure:
                logger.warning("MQTT connection refused: %s", reason_code)
                self._connected = False
            else:
                logger.info("MQTT connected (reason_code=%s).", reason_code)
                self._connected = True
                # Re-establish all registered subscriptions after (re-)connect.
                for sub_topic in self._subscriptions:
                    try:
                        _client.subscribe(sub_topic)
                        logger.debug("MQTT re-subscribed to topic=%s.", sub_topic)
                    except Exception:
                        logger.exception(
                            "MQTT re-subscribe failed (topic=%s).", sub_topic
                        )

        # VERSION2 on_disconnect signature:
        # (client, userdata, disconnect_flags, reason_code, properties)
        def _on_disconnect(
            _client: mqtt.Client,
            _userdata: object,
            _disconnect_flags: mqtt.DisconnectFlags,
            reason_code: mqtt.ReasonCode,
            _properties: mqtt.Properties | None,
        ) -> None:
            self._connected = False
            if reason_code.is_failure:
                logger.warning("MQTT disconnected unexpectedly (reason_code=%s).", reason_code)
            else:
                logger.info("MQTT disconnected cleanly.")

        # VERSION2 on_message signature: (client, userdata, message)
        def _on_message(
            _client: mqtt.Client,
            _userdata: object,
            message: mqtt.MQTTMessage,
        ) -> None:
            handler = self._subscriptions.get(message.topic)
            if handler is None:
                logger.debug(
                    "MQTT on_message: no handler for topic=%s.", message.topic
                )
                return
            try:
                handler(message.payload)
            except Exception:
                logger.exception(
                    "MQTT message handler raised for topic=%s (swallowed).",
                    message.topic,
                )

        client.on_connect = _on_connect
        client.on_disconnect = _on_disconnect
        client.on_message = _on_message

        # TLS
        if settings.mqtt_tls_enabled:
            try:
                client.tls_set()
            except Exception:
                logger.exception("MQTT TLS setup failed.")
                return

        # Credentials — never log the password
        if settings.mqtt_username:
            client.username_pw_set(
                username=settings.mqtt_username,
                password=settings.mqtt_password or None,
            )

        # Start background loop thread *before* connect so paho can handle
        # the connection handshake asynchronously.
        client.loop_start()

        try:
            client.connect(
                host=settings.mqtt_broker_host,
                port=settings.mqtt_broker_port,
                keepalive=60,
            )
        except Exception:
            # Log without leaking the password
            logger.exception(
                "MQTT connect failed (host=%s, port=%s). Stopping loop.",
                settings.mqtt_broker_host,
                settings.mqtt_broker_port,
            )
            try:
                client.loop_stop()
            except Exception:
                pass
            return

        self._client = client
        logger.info(
            "MQTT client started (host=%s, port=%s).",
            settings.mqtt_broker_host,
            settings.mqtt_broker_port,
        )

    def _stop_client(self) -> None:
        """Disconnect and stop the background loop.  Idempotent."""
        client = self._client
        if client is None:
            return
        self._client = None
        self._connected = False
        try:
            client.disconnect()
        except Exception:
            logger.debug("MQTT disconnect raised (ignoring).", exc_info=True)
        try:
            client.loop_stop()
        except Exception:
            logger.debug("MQTT loop_stop raised (ignoring).", exc_info=True)
        logger.info("MQTT client stopped.")

    def _stop_source_client(self, source_id: int) -> None:
        """Detach then stop a source client without blocking callback bookkeeping.

        Callers hold ``_source_lifecycle_lock``.  The first phase makes the
        generation unreachable while holding ``_lock``.  Paho operations and
        the in-flight wait are deliberately outside that lock: loop_stop()
        joins paho's network thread, and an active callback needs ``_lock`` to
        release its permit in ``_on_message``'s finally block.
        """
        with self._lock:
            state = self._source_states.pop(source_id, None)
            client = self._source_clients.pop(source_id, None)
            self._source_subscriptions.pop(source_id, None)
            self._source_connected.discard(source_id)
            # Invalidate callbacks even when there was no successfully
            # installed client (for example after a failed replacement).
            self._source_generations.pop(source_id, None)
        if client is not None:
            try:
                client.disconnect()
            except Exception:
                logger.debug("DSMR MQTT disconnect raised (source_id=%s)", source_id, exc_info=True)
            try:
                client.loop_stop()
            except Exception:
                logger.debug("DSMR MQTT loop_stop raised (source_id=%s)", source_id, exc_info=True)
        if state is not None:
            with self._lock:
                while state.in_flight:
                    self._source_idle.wait()

    def _is_current_source_client(
        self, source_id: int, generation: int, client: mqtt.Client
    ) -> bool:
        """Check callback ownership while ``_lock`` is held."""
        return (
            self._source_generations.get(source_id) == generation
            and self._source_clients.get(source_id) is client
        )

    @staticmethod
    def _report_source_state(
        state_handler: Callable[[str], None] | None, state: str, source_id: int
    ) -> None:
        """Invoke an optional health callback without exposing connection credentials."""
        if state_handler is None:
            return
        try:
            state_handler(state)
        except Exception:
            logger.exception("DSMR MQTT source state update failed for source_id=%s", source_id)


# ---------------------------------------------------------------------------
# Module-level singleton — shared across lifespan and route handlers
# ---------------------------------------------------------------------------

mqtt_manager = MqttManager()