"""Thin paho-mqtt wrapper.

Importing this module does not connect to a broker. Network access starts only
when MqttClient.connect() is called.
"""

from __future__ import annotations

from collections.abc import Callable
import threading

import paho.mqtt.client as mqtt


MessageCallback = Callable[[str, bytes, mqtt.MQTTMessage], None]
OnConnectHook = Callable[[], None]


class MqttClient:
    def __init__(self, client_id: str | None = None) -> None:
        self._connected = threading.Event()
        self._state_lock = threading.Lock()
        self._subscriptions: dict[str, tuple[MessageCallback, int]] = {}
        self._on_connect_hooks: list[OnConnectHook] = []
        self._client = mqtt.Client(
            callback_api_version=mqtt.CallbackAPIVersion.VERSION2,
            client_id=client_id or "",
        )
        self._client.on_connect = self._on_connect

    def _on_connect(
        self,
        client: mqtt.Client,
        userdata: object,
        flags: mqtt.ConnectFlags,
        reason_code: mqtt.ReasonCode,
        properties: mqtt.Properties | None,
    ) -> None:
        if reason_code.is_failure:
            print(f"MQTT connection failed: {reason_code}", flush=True)
            return

        with self._state_lock:
            subscriptions = tuple(
                (topic, qos) for topic, (_, qos) in self._subscriptions.items()
            )
            hooks = tuple(self._on_connect_hooks)

        for topic, qos in subscriptions:
            self._client.subscribe(topic, qos=qos)

        for hook in hooks:
            try:
                hook()
            except Exception as exc:
                print(f"MQTT on-connect hook failed: {exc}", flush=True)

        self._connected.set()

    def add_on_connect_hook(self, fn: OnConnectHook) -> None:
        with self._state_lock:
            self._on_connect_hooks.append(fn)

    def set_last_will(
        self,
        topic: str,
        payload: str | bytes,
        qos: int = 1,
        retain: bool = True,
    ) -> None:
        self._client.will_set(topic, payload=payload, qos=qos, retain=retain)

    def connect(
        self,
        host: str,
        port: int,
        username: str | None = None,
        password: str | None = None,
        tls: bool = False,
    ) -> None:
        if username is not None:
            self._client.username_pw_set(username, password)
        if tls:
            self._client.tls_set()
        self._client.connect(host, port)

    def publish(
        self,
        topic: str,
        payload: str | bytes,
        qos: int = 1,
        retain: bool = False,
    ) -> mqtt.MQTTMessageInfo:
        return self._client.publish(topic, payload=payload, qos=qos, retain=retain)

    def subscribe(self, topic: str, callback: MessageCallback, qos: int = 1) -> None:
        def _handle_message(
            client: mqtt.Client,
            userdata: object,
            message: mqtt.MQTTMessage,
        ) -> None:
            callback(message.topic, message.payload, message)

        self._client.message_callback_add(topic, _handle_message)
        with self._state_lock:
            self._subscriptions[topic] = (callback, qos)
        self._client.subscribe(topic, qos=qos)

    def loop_start(self) -> None:
        self._client.loop_start()

    def loop_stop(self) -> None:
        self._client.loop_stop()

    def wait_connected(self, timeout: float = 10.0) -> bool:
        return self._connected.wait(timeout)

    def disconnect(self) -> None:
        self._client.disconnect()
