"""Per-broker MQTT runtime — maintains a persistent connection, topic cache, and pub/sub.

Each MQTTSource entity gets its own MQTTRuntime instance, managed by MQTTManager.
Refactored from the former singleton MQTTService.
"""

import asyncio
import json
from typing import Callable, Dict, Optional, Set

import aiomqtt

from wled_controller.storage.mqtt_source import MQTTSource
from wled_controller.utils import get_logger

logger = get_logger(__name__)


class MQTTRuntime:
    """Persistent MQTT broker connection for a single MQTTSource.

    Features:
    - Publish messages (retained or transient)
    - Subscribe to topics with callback dispatch
    - Topic value cache for synchronous reads (automation rule evaluation)
    - Auto-reconnect loop
    - Birth / will messages for online status
    """

    _RECONNECT_DELAY = 5.0  # seconds between reconnect attempts

    def __init__(self, source: MQTTSource) -> None:
        self._source_id = source.id
        self._broker_host = source.broker_host
        self._broker_port = source.broker_port
        self._username = source.username
        self._password = source.password
        self._client_id = source.client_id
        self._base_topic = source.base_topic

        self._client: Optional[aiomqtt.Client] = None
        self._task: Optional[asyncio.Task] = None
        self._connected = False

        # Subscription management
        self._subscriptions: Dict[str, Set[Callable]] = {}  # topic -> set of callbacks
        self._topic_cache: Dict[str, str] = {}  # topic -> last payload string

        # Pending publishes queued while disconnected
        self._publish_queue: asyncio.Queue = asyncio.Queue(maxsize=1000)

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

    @property
    def source_id(self) -> str:
        return self._source_id

    @property
    def base_topic(self) -> str:
        return self._base_topic

    async def start(self) -> None:
        """Start the MQTT connection loop."""
        if self._task is not None:
            return
        self._task = asyncio.create_task(self._connection_loop())
        logger.info(
            "MQTT runtime started: %s -> %s:%d",
            self._source_id,
            self._broker_host,
            self._broker_port,
        )

    async def stop(self) -> None:
        """Stop the MQTT connection."""
        if self._task is None:
            return
        self._task.cancel()
        try:
            await self._task
        except asyncio.CancelledError:
            logger.debug("MQTT runtime task cancelled: %s", self._source_id)
        self._task = None
        self._connected = False
        logger.info("MQTT runtime stopped: %s", self._source_id)

    def update_config(self, source: MQTTSource) -> None:
        """Hot-update config. Connection will use new values on next reconnect."""
        self._broker_host = source.broker_host
        self._broker_port = source.broker_port
        self._username = source.username
        self._password = source.password
        self._client_id = source.client_id
        self._base_topic = source.base_topic

    async def publish(self, topic: str, payload: str, retain: bool = False, qos: int = 0) -> None:
        """Publish a message. If disconnected, queue for later delivery."""
        if self._connected and self._client:
            try:
                await self._client.publish(topic, payload, retain=retain, qos=qos)
                return
            except Exception as e:
                logger.warning("MQTT publish failed (%s/%s): %s", self._source_id, topic, e)
        # Queue for retry
        try:
            self._publish_queue.put_nowait((topic, payload, retain, qos))
        except asyncio.QueueFull:
            logger.warning(
                "MQTT publish queue full (%s), dropping message for topic %s",
                self._source_id,
                topic,
            )

    async def subscribe(self, topic: str, callback: Callable) -> None:
        """Subscribe to a topic. Callback receives (topic: str, payload: str)."""
        if topic not in self._subscriptions:
            self._subscriptions[topic] = set()
        self._subscriptions[topic].add(callback)

        # Subscribe on the live client if connected
        if self._connected and self._client:
            try:
                await self._client.subscribe(topic)
            except Exception as e:
                logger.warning("MQTT subscribe failed (%s/%s): %s", self._source_id, topic, e)

    def get_last_value(self, topic: str) -> Optional[str]:
        """Get cached last value for a topic (synchronous — for automation evaluation)."""
        return self._topic_cache.get(topic)

    def get_subscribed_topics(self) -> int:
        """Return number of subscribed topics."""
        return len(self._subscriptions)

    async def _connection_loop(self) -> None:
        """Persistent connection loop with auto-reconnect."""
        while True:
            will_topic = f"{self._base_topic}/status"
            try:
                async with aiomqtt.Client(
                    hostname=self._broker_host,
                    port=self._broker_port,
                    username=self._username or None,
                    password=self._password or None,
                    identifier=self._client_id,
                    will=aiomqtt.Will(
                        topic=will_topic,
                        payload="offline",
                        retain=True,
                    ),
                ) as client:
                    self._client = client
                    self._connected = True
                    logger.info(
                        "MQTT connected: %s -> %s:%d",
                        self._source_id,
                        self._broker_host,
                        self._broker_port,
                    )

                    # Publish birth message
                    await client.publish(will_topic, "online", retain=True)

                    # Re-subscribe to all registered topics
                    for topic in self._subscriptions:
                        await client.subscribe(topic)

                    # Drain pending publishes
                    while not self._publish_queue.empty():
                        try:
                            t, p, r, q = self._publish_queue.get_nowait()
                            await client.publish(t, p, retain=r, qos=q)
                        except Exception:
                            break

                    # Message receive loop
                    async for msg in client.messages:
                        topic_str = str(msg.topic)
                        payload_str = (
                            msg.payload.decode("utf-8", errors="replace") if msg.payload else ""
                        )
                        self._topic_cache[topic_str] = payload_str

                        # Dispatch to callbacks
                        for sub_topic, callbacks in self._subscriptions.items():
                            if aiomqtt.Topic(sub_topic).matches(msg.topic):
                                for cb in callbacks:
                                    try:
                                        if asyncio.iscoroutinefunction(cb):
                                            asyncio.create_task(cb(topic_str, payload_str))
                                        else:
                                            cb(topic_str, payload_str)
                                    except Exception as e:
                                        logger.error(
                                            "MQTT callback error (%s/%s): %s",
                                            self._source_id,
                                            topic_str,
                                            e,
                                        )

            except asyncio.CancelledError:
                break
            except Exception as e:
                self._connected = False
                self._client = None
                logger.warning(
                    "MQTT connection lost (%s): %s. Reconnecting in %.0fs...",
                    self._source_id,
                    e,
                    self._RECONNECT_DELAY,
                )
                await asyncio.sleep(self._RECONNECT_DELAY)

    # ===== State exposure helpers =====

    async def publish_target_state(self, target_id: str, state: dict) -> None:
        """Publish target state to MQTT (called from event handler)."""
        topic = f"{self._base_topic}/target/{target_id}/state"
        await self.publish(topic, json.dumps(state), retain=True)

    async def publish_automation_state(self, automation_id: str, action: str) -> None:
        """Publish automation state change to MQTT."""
        topic = f"{self._base_topic}/automation/{automation_id}/state"
        await self.publish(topic, json.dumps({"action": action}), retain=True)
