"""MQTT connection manager using aiomqtt.

Connects to EMQX for WebRTC signaling, telemetry reporting, and command
delivery. Completely independent from the ACP (Bridge) channel.
"""

from __future__ import annotations

import asyncio
import json
import logging
from typing import Any, Callable, Awaitable

import aiomqtt

logger = logging.getLogger(__name__)

MessageCallback = Callable[[str, dict[str, Any]], Awaitable[None]]

_MAX_RECONNECT_DELAY = 60.0
_BASE_RECONNECT_DELAY = 1.0


class MQTTManager:
    """Manages the MQTT connection to EMQX for signaling and telemetry."""

    def __init__(
        self,
        device_id: str,
        mqtt_host: str,
        mqtt_port: int = 8883,
        mqtt_username: str | None = None,
        mqtt_password: str | None = None,
        use_tls: bool = True,
        auto_reconnect: bool = True,
    ) -> None:
        self._device_id = device_id
        self._mqtt_host = mqtt_host
        self._mqtt_port = mqtt_port
        self._mqtt_username = mqtt_username
        self._mqtt_password = mqtt_password
        self._use_tls = use_tls
        self._auto_reconnect = auto_reconnect

        self._client: aiomqtt.Client | None = None
        self._connected = False
        self._shutting_down = False
        self._handlers: dict[str, list[MessageCallback]] = {}
        self._listen_task: asyncio.Task[None] | None = None

    @property
    def device_topic_prefix(self) -> str:
        return f"devices/{self._device_id}"

    @property
    def connected(self) -> bool:
        return self._connected

    def on_topic(self, suffix: str, callback: MessageCallback) -> None:
        """Register a handler for messages on ``devices/{deviceId}/{suffix}``."""
        full_topic = f"{self.device_topic_prefix}/{suffix}"
        self._handlers.setdefault(full_topic, []).append(callback)

    async def connect(self) -> None:
        """Connect to the MQTT broker and start listening for messages."""
        self._shutting_down = False
        await self._do_connect()
        self._listen_task = asyncio.create_task(self._run_loop())

    async def _do_connect(self) -> None:
        self._client = aiomqtt.Client(
            hostname=self._mqtt_host,
            port=self._mqtt_port,
            identifier=f"device-{self._device_id}",
            username=self._mqtt_username,
            password=self._mqtt_password,
            tls_params=aiomqtt.TLSParameters() if self._use_tls else None,
        )
        await self._client.__aenter__()
        self._connected = True

        for topic in self._handlers:
            await self._client.subscribe(topic)
            logger.debug("Subscribed to %s", topic)

        logger.info("MQTT connected to %s:%d", self._mqtt_host, self._mqtt_port)

    async def disconnect(self) -> None:
        """Disconnect from the MQTT broker."""
        self._shutting_down = True

        if self._listen_task and not self._listen_task.done():
            self._listen_task.cancel()
            try:
                await self._listen_task
            except asyncio.CancelledError:
                pass
            self._listen_task = None

        await self._close_client()

    async def _close_client(self) -> None:
        if self._client:
            try:
                await self._client.__aexit__(None, None, None)
            except Exception:
                pass
            self._client = None
        self._connected = False

    async def publish(self, suffix: str, payload: dict[str, Any]) -> None:
        """Publish a JSON message to ``devices/{deviceId}/{suffix}``."""
        if not self._client:
            raise RuntimeError("MQTT not connected")
        topic = f"{self.device_topic_prefix}/{suffix}"
        await self._client.publish(topic, json.dumps(payload).encode())

    async def _run_loop(self) -> None:
        """Main loop with auto-reconnect. Delegates to _listen_loop for message dispatch."""
        delay = _BASE_RECONNECT_DELAY
        while not self._shutting_down:
            try:
                await self._listen_loop()
            except asyncio.CancelledError:
                return
            except Exception:
                logger.exception("MQTT connection lost")

            self._connected = False
            await self._close_client()

            if self._shutting_down or not self._auto_reconnect:
                return

            logger.info("MQTT reconnecting in %.1fs ...", delay)
            await asyncio.sleep(delay)
            delay = min(delay * 2, _MAX_RECONNECT_DELAY)

            try:
                await self._do_connect()
                delay = _BASE_RECONNECT_DELAY
            except asyncio.CancelledError:
                return
            except Exception:
                logger.exception("MQTT reconnect failed")

    async def _listen_loop(self) -> None:
        """Listen for incoming MQTT messages and dispatch to handlers."""
        if not self._client:
            return
        async for message in self._client.messages:
            topic = str(message.topic)
            raw = message.payload
            if isinstance(raw, (bytes, bytearray)):
                raw = raw.decode("utf-8", errors="replace")
            try:
                payload = json.loads(raw)
            except (json.JSONDecodeError, TypeError):
                logger.warning("Non-JSON MQTT message on %s", topic)
                continue

            handlers = self._handlers.get(topic, [])
            for handler in handlers:
                try:
                    await handler(topic, payload)
                except Exception:
                    logger.exception("MQTT handler error on %s", topic)
