"""MQTT runtime manager — ref-counted pool of MQTT broker connections.

Follows the HomeAssistantManager pattern: multiple consumers (MQTT devices,
automation rules, state publishing) sharing the same broker connection per
MQTTSource instance.
"""

import asyncio
from typing import Any, Dict, List, Optional

from wled_controller.core.mqtt.mqtt_runtime import MQTTRuntime
from wled_controller.storage.mqtt_source_store import MQTTSourceStore
from wled_controller.utils import get_logger

logger = get_logger(__name__)


class MQTTManager:
    """Ref-counted pool of MQTT runtimes.

    Each MQTT source gets at most one runtime (broker connection).
    Multiple consumers share the same runtime via acquire/release.
    """

    def __init__(self, store: MQTTSourceStore) -> None:
        self._store = store
        # source_id -> (runtime, ref_count)
        self._runtimes: Dict[str, tuple] = {}
        self._lock = asyncio.Lock()

    async def acquire(self, source_id: str) -> MQTTRuntime:
        """Get or create a runtime for the given MQTT source. Increments ref count."""
        async with self._lock:
            if source_id in self._runtimes:
                runtime, count = self._runtimes[source_id]
                self._runtimes[source_id] = (runtime, count + 1)
                return runtime

            source = self._store.get(source_id)
            runtime = MQTTRuntime(source)
            await runtime.start()
            self._runtimes[source_id] = (runtime, 1)
            return runtime

    async def release(self, source_id: str) -> None:
        """Decrement ref count; stop runtime when it reaches zero."""
        async with self._lock:
            if source_id not in self._runtimes:
                return
            runtime, count = self._runtimes[source_id]
            if count <= 1:
                await runtime.stop()
                del self._runtimes[source_id]
            else:
                self._runtimes[source_id] = (runtime, count - 1)

    def get_last_value(self, source_id: str, topic: str) -> Optional[str]:
        """Get cached topic value from a running runtime (synchronous)."""
        entry = self._runtimes.get(source_id)
        if entry is None:
            return None
        runtime, _count = entry
        return runtime.get_last_value(topic)

    def get_runtime(self, source_id: str) -> Optional[MQTTRuntime]:
        """Get a running runtime without changing ref count (for read-only access)."""
        entry = self._runtimes.get(source_id)
        if entry is None:
            return None
        runtime, _count = entry
        return runtime

    async def ensure_runtime(self, source_id: str) -> MQTTRuntime:
        """Get or create a runtime (for API endpoints that need a connection)."""
        async with self._lock:
            if source_id in self._runtimes:
                runtime, _count = self._runtimes[source_id]
                return runtime

        source = self._store.get(source_id)
        runtime = MQTTRuntime(source)
        await runtime.start()
        async with self._lock:
            if source_id not in self._runtimes:
                self._runtimes[source_id] = (runtime, 0)
            else:
                await runtime.stop()
                runtime, _count = self._runtimes[source_id]
        return runtime

    async def update_source(self, source_id: str) -> None:
        """Hot-update runtime config when the source is edited."""
        entry = self._runtimes.get(source_id)
        if entry is None:
            return
        runtime, _count = entry
        try:
            source = self._store.get(source_id)
            runtime.update_config(source)
        except Exception as e:
            logger.warning("Failed to update MQTT runtime %s: %s", source_id, e)

    def get_connection_status(self) -> List[Dict[str, Any]]:
        """Get status of all active MQTT connections (for dashboard indicators)."""
        result = []
        for source_id, (runtime, ref_count) in self._runtimes.items():
            try:
                source = self._store.get(source_id)
                name = source.name
            except Exception:
                name = source_id
            result.append(
                {
                    "source_id": source_id,
                    "name": name,
                    "connected": runtime.is_connected,
                    "ref_count": ref_count,
                    "subscribed_topics": runtime.get_subscribed_topics(),
                }
            )
        return result

    def get_all_sources_status(self) -> List[Dict[str, Any]]:
        """Get status for ALL MQTT sources (connected or not), for the dashboard."""
        result = []
        for source in self._store.get_all():
            entry = self._runtimes.get(source.id)
            if entry:
                runtime, ref_count = entry
                result.append(
                    {
                        "source_id": source.id,
                        "name": source.name,
                        "connected": runtime.is_connected,
                        "ref_count": ref_count,
                        "broker": f"{source.broker_host}:{source.broker_port}",
                    }
                )
            else:
                result.append(
                    {
                        "source_id": source.id,
                        "name": source.name,
                        "connected": False,
                        "ref_count": 0,
                        "broker": f"{source.broker_host}:{source.broker_port}",
                    }
                )
        return result

    def get_first_runtime(self) -> Optional[MQTTRuntime]:
        """Get the first available connected runtime (fallback for unlinked consumers)."""
        for _source_id, (runtime, _count) in self._runtimes.items():
            if runtime.is_connected:
                return runtime
        # Return any runtime even if disconnected
        for _source_id, (runtime, _count) in self._runtimes.items():
            return runtime
        return None

    async def shutdown(self) -> None:
        """Stop all runtimes."""
        async with self._lock:
            for source_id, (runtime, _count) in list(self._runtimes.items()):
                await runtime.stop()
            self._runtimes.clear()
        logger.info("MQTT manager shut down")
