import asyncio
import logging
import ssl
from typing import Awaitable, Callable, Optional, Tuple, List

import paho.mqtt.client as mqtt

from gateway_transport.mqtt.MqttClientProtocol import MqttClientProtocol

logger = logging.getLogger(__name__)


class PahoMqttClient(MqttClientProtocol):
    """
    MQTT transport adapter built on paho-mqtt.

    Implements MqttClientProtocol so inbound and outbound adapters
    (DeltaGrid listener + DeltaGridClient) can use it.
    """

    def __init__(
        self,
        host: str,
        port: int = 8883,
        username: Optional[str] = None,
        password: Optional[str] = None,
        cafile: Optional[str] = None,
        clientId: Optional[str] = None,
        keepalive: int = 60,
        useTls: bool = True,
    ) -> None:
        self._host = host
        self._port = port
        self._username = username
        self._password = password
        self._cafile = cafile
        self._keepalive = keepalive
        self._useTls = useTls

        # Asyncio loop and callback coming from the hexagon layer
        self._loop: Optional[asyncio.AbstractEventLoop] = None
        self._messageCallback: Optional[Callable[[str, bytes], Awaitable[None]]] = None

        # Subscriptions so we can resubscribe after reconnect
        self._subscriptions: List[Tuple[str, int]] = []

        # Initialize paho client
        self._client = mqtt.Client(
            client_id=clientId or "",
            protocol=mqtt.MQTTv5,  # adjust if you need a different version
        )

        if self._username:
            self._client.username_pw_set(self._username, self._password)

        if self._useTls:
            self._client.tls_set(
                ca_certs=self._cafile,
                cert_reqs=ssl.CERT_REQUIRED,
                tls_version=ssl.PROTOCOL_TLS,
            )
            self._client.tls_insecure_set(False)

        self._client.on_connect = self._on_connect
        self._client.on_message = self._on_message
        self._client.on_disconnect = self._on_disconnect
        self._connectedEvent = asyncio.Event()

    # -------------------------------------------------------------------------
    # MqttClientProtocol implementation
    # -------------------------------------------------------------------------
    async def subscribe(
        self,
        topic: str,
        callback: Callable[[str, bytes], Awaitable[None]],
    ) -> None:
        """
        Called by the DeltaGrid listener to subscribe to a topic pattern.
        """
        # Store the event loop and callback from the async context
        self._loop = asyncio.get_running_loop()
        self._messageCallback = callback
        entry = (topic, 0)
        if entry not in self._subscriptions:
            self._subscriptions.append(entry)

        # If we're already connected, subscribe immediately
        if self._client.is_connected():
            logger.info("Subscribing to topic %s", topic)
            result, _ = self._client.subscribe(topic, qos=0)
            if result != mqtt.MQTT_ERR_SUCCESS:
                logger.error("Failed to subscribe to topic %s (code=%s)", topic, result)

    async def publish(
        self,
        topic: str,
        payload: bytes,
        qos: int = 0,
        retain: bool = False,
    ) -> None:
        """
        Used by the DeltaGridClient (outbound).
        """
        # publish is thread-safe because paho runs its own loop thread
        result = self._client.publish(topic, payload=payload, qos=qos, retain=retain)
        if result.rc != mqtt.MQTT_ERR_SUCCESS:
            logger.error(
                "Failed to publish to topic %s (code=%s)", topic, result.rc
            )

    async def runForever(self) -> None:
        """
        Connect and keep the MQTT loop running until cancelled.
        Uses paho's internal loop thread (loop_start/loop_stop) so we can stop cleanly.
        """
        self._loop = asyncio.get_running_loop()

        logger.info(
            "Connecting to MQTT broker %s:%s (TLS=%s)",
            self._host,
            self._port,
            self._useTls,
        )

        # connect() is blocking but usually fast; if you want, move it to to_thread.
        self._client.connect(self._host, self._port, keepalive=self._keepalive)

        # start paho network loop in its own thread
        self._client.loop_start()

        try:
            # block until cancelled
            await asyncio.Event().wait()
        except asyncio.CancelledError:
            logger.info("Stopping MQTT client")
            raise
        finally:
            # ensure we stop the paho loop thread
            try:
                self._client.disconnect()
            except Exception:  # noqa: BLE001
                logger.exception("MQTT disconnect failed")
            finally:
                # loop_stop(force=True) stops the thread even if disconnect is slow
                self._client.loop_stop(force=True)

            if self._loop is not None:
                self._loop.call_soon_threadsafe(self._connectedEvent.clear)

    # -------------------------------------------------------------------------
    # Paho callbacks (run in their own thread)
    # -------------------------------------------------------------------------
    def _on_connect(self, client, userdata, flags, reason_code, properties=None):
        if reason_code != mqtt.MQTT_ERR_SUCCESS:
            logger.error("MQTT connect failed: %s", reason_code)
            return

        logger.info("Connected to MQTT broker %s:%s", self._host, self._port)

        if self._loop is not None:
            self._loop.call_soon_threadsafe(self._connectedEvent.set)

        # After connect/reconnect, resubscribe to all topics
        for topic, qos in self._subscriptions:
            logger.info("Subscribing to topic %s", topic)
            result, _ = self._client.subscribe(topic, qos=qos)
            if result != mqtt.MQTT_ERR_SUCCESS:
                logger.error(
                    "Failed to subscribe to topic %s (code=%s)", topic, result
                )

    def _on_disconnect(self, client, userdata, reason_code, properties=None):
        logger.warning("Disconnected from MQTT broker (code=%s)", reason_code)

        if self._loop is not None:
            self._loop.call_soon_threadsafe(self._connectedEvent.clear)

    def _on_message(self, client, userdata, msg):
        """
        Called by paho in its own thread.
        We hand processing back into the asyncio event loop.
        """
        if self._loop is None or self._messageCallback is None:
            logger.warning("Received MQTT message but no callback is registered")
            return

        topic = msg.topic
        payload = msg.payload

        async def _call():
            try:
                await self._messageCallback(topic, payload)
            except Exception:  # noqa: BLE001
                logger.exception("Error while handling MQTT message on %s", topic)

        # Safely jump back into the asyncio loop
        self._loop.call_soon_threadsafe(asyncio.create_task, _call())

    async def waitConnected(self, timeout: float | None = 10.0) -> None:
        if timeout is None:
            await self._connectedEvent.wait()
            return
        await asyncio.wait_for(self._connectedEvent.wait(), timeout=timeout)