"""Async MQTT client wrapper used by Rex voice nodes.

The implementation is tailored for TLS-enabled brokers with automatic
reconnection, payload validation, and a watchdog heartbeat.
"""

from __future__ import annotations

import asyncio
import inspect
import json
import logging
import ssl
import time
from collections.abc import Awaitable, Callable, Iterable
from contextlib import suppress
from dataclasses import dataclass, field
from datetime import UTC, datetime, timedelta
from pathlib import Path
from typing import TYPE_CHECKING, Any
from uuid import uuid4

from rex.assistant_errors import ConfigurationError
from rex.config import settings

if TYPE_CHECKING:
    from aiomqtt import Client, Message, MqttError
else:
    try:
        # Try new package name first (asyncio-mqtt was renamed to aiomqtt)
        from aiomqtt import Client, Message, MqttError
    except ImportError:
        try:
            # Fall back to old package name for backwards compatibility
            from asyncio_mqtt import Client, Message, MqttError
        except ImportError:  # pragma: no cover - optional dependency
            Client = None
            Message = None
            MqttError = Exception

logger = logging.getLogger(__name__)

MessageHandler = Callable[[str, dict[str, Any]], Awaitable[None] | None]


class PayloadValidationError(RuntimeError):
    """Raised when a payload does not meet the expected schema."""


def _topic_matches(subscription: str, topic: str) -> bool:
    """Return True if an MQTT topic matches a subscription filter."""
    sub_levels = subscription.split("/")
    topic_levels = topic.split("/")

    for idx, level in enumerate(sub_levels):
        if level == "#":
            return True
        if idx >= len(topic_levels):
            return False
        if level in {"+", topic_levels[idx]}:
            continue
        return False

    return len(topic_levels) == len(sub_levels)


def _now_iso() -> str:
    """Return the current UTC timestamp in ISO-8601 format."""
    return datetime.now(tz=UTC).isoformat()


def _resolve_optional_path(path: str | None) -> Path | None:
    if not path:
        return None
    resolved = Path(path).expanduser()
    return resolved if resolved.exists() else None


@dataclass
class Subscription:
    topic: str
    qos: int = 1
    handlers: list[MessageHandler] = field(default_factory=list)


class RexMQTTClient:
    """High-level MQTT client with TLS, watchdog, and payload validation."""

    DEFAULT_SUBSCRIPTIONS: tuple[str, ...] = ("rex/audio_in", "rex/audio_out", "rex/ha_cmd")

    def __init__(
        self,
        *,
        broker: str | None = None,
        port: int | None = None,
        tls_enabled: bool | None = None,
        username: str | None = None,
        password: str | None = None,
        client_id: str | None = None,
        keepalive: int | None = None,
        tls_ca: str | None = None,
        tls_cert: str | None = None,
        tls_key: str | None = None,
        tls_insecure: bool | None = None,
        watchdog_interval: int | None = None,
        watchdog_timeout: int | None = None,
        node_id: str | None = None,
        health_topic: str = "rex/health/core",
        required_fields: Iterable[str] = ("node_id", "timestamp"),
        reconnect_interval: float = 5.0,
        reconnect_interval_max: float = 60.0,
    ) -> None:
        cfg = settings
        self._broker = broker or cfg.mqtt_broker
        self._port = int(port or cfg.mqtt_port)
        self._tls_enabled = bool(cfg.mqtt_tls if tls_enabled is None else tls_enabled)
        self._username = username if username is not None else cfg.mqtt_username
        self._password = password if password is not None else cfg.mqtt_password
        self._client_id = (client_id or cfg.mqtt_client_id or f"rex-core-{uuid4().hex[:8]}").strip()
        self._keepalive = int(keepalive or cfg.mqtt_keepalive)
        self._tls_ca = tls_ca if tls_ca is not None else cfg.mqtt_tls_ca
        self._tls_cert = tls_cert if tls_cert is not None else cfg.mqtt_tls_cert
        self._tls_key = tls_key if tls_key is not None else cfg.mqtt_tls_key
        self._tls_insecure = bool(cfg.mqtt_tls_insecure if tls_insecure is None else tls_insecure)
        self._watchdog_interval = int(watchdog_interval or cfg.mqtt_watchdog_interval)
        self._watchdog_timeout = int(watchdog_timeout or cfg.mqtt_watchdog_timeout)
        self._node_id = node_id or cfg.mqtt_node_id or "rex_core"
        self._health_topic = health_topic
        self._required_fields = tuple(required_fields)
        self._reconnect_interval = reconnect_interval
        self._reconnect_interval_max = reconnect_interval_max

        self._subscriptions: dict[str, Subscription] = {
            topic: Subscription(topic=topic) for topic in self.DEFAULT_SUBSCRIPTIONS
        }

        self._client: Client | None = None
        self._connection_task: asyncio.Task[None] | None = None
        self._watchdog_task: asyncio.Task[None] | None = None
        self._stop_event = asyncio.Event()
        self._connected_event = asyncio.Event()
        self._last_message_at = time.monotonic()
        if not self._broker:
            raise ConfigurationError("MQTT broker host is required (set REX_MQTT_BROKER).")

    # --------------------------------------------------------------------- #
    # Public API
    # --------------------------------------------------------------------- #

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

    async def start(self, *, wait_connected: bool = True, timeout: float | None = 20.0) -> None:
        """Start the MQTT connection loop and optional watchdog."""
        if self._connection_task is None or self._connection_task.done():
            self._stop_event.clear()
            self._connection_task = asyncio.create_task(
                self._connection_loop(), name="rex-mqtt-connection"
            )
        if self._watchdog_task is None or self._watchdog_task.done():
            self._watchdog_task = asyncio.create_task(
                self._watchdog_loop(), name="rex-mqtt-watchdog"
            )

        if wait_connected:
            await self.wait_connected(timeout=timeout)

    async def stop(self) -> None:
        """Stop connection loop and close the MQTT client."""
        self._stop_event.set()

        for task in (self._watchdog_task, self._connection_task):
            if task is None:
                continue
            task.cancel()
            with suppress(asyncio.CancelledError):
                await task

        self._watchdog_task = None
        self._connection_task = None

        client = self._client
        self._client = None
        if client is not None:
            with suppress(MqttError):
                await client.disconnect()

        self._connected_event.clear()

    async def wait_connected(self, *, timeout: float | None = None) -> None:
        """Block until the client is connected or timeout expires."""
        if timeout is None:
            await self._connected_event.wait()
            return
        try:
            await asyncio.wait_for(self._connected_event.wait(), timeout)
        except TimeoutError as exc:
            raise TimeoutError("MQTT connection timed out.") from exc

    async def publish(
        self,
        topic: str,
        payload: dict[str, Any] | str | bytes,
        *,
        qos: int = 1,
        retain: bool = False,
        ensure_base_fields: bool = True,
    ) -> None:
        """Publish a message to the broker."""
        client = self._client
        if client is None or not self.is_connected:
            raise RuntimeError("MQTT client is not connected.")

        if isinstance(payload, dict):
            message = dict(payload)
            if ensure_base_fields:
                message.setdefault("node_id", self._node_id)
                message.setdefault("timestamp", _now_iso())
            self._validate_payload(message, context=f"publish:{topic}", check_timestamp=False)
            try:
                payload_bytes = json.dumps(message, separators=(",", ":")).encode("utf-8")
            except (TypeError, ValueError) as exc:
                raise PayloadValidationError(
                    f"[publish:{topic}] Payload is not JSON serializable."
                ) from exc
        elif isinstance(payload, str):
            payload_bytes = payload.encode("utf-8")
        else:
            payload_bytes = payload

        try:
            await client.publish(topic, payload_bytes, qos=qos, retain=retain)
        except MqttError:
            self._connected_event.clear()
            raise

    async def add_subscription(
        self,
        topic: str,
        handler: MessageHandler,
        *,
        qos: int = 1,
        immediate_sync: bool = True,
    ) -> None:
        """Register a handler for a topic filter and ensure subscription."""
        subscription = self._subscriptions.get(topic)
        if subscription is None:
            subscription = Subscription(topic=topic, qos=qos)
            self._subscriptions[topic] = subscription
        else:
            subscription.qos = max(subscription.qos, qos)
        subscription.handlers.append(handler)

        if immediate_sync and self._client is not None and self.is_connected:
            await self._client.subscribe(topic, qos=subscription.qos)

    # ------------------------------------------------------------------ #
    # Internal helpers
    # ------------------------------------------------------------------ #

    async def _connection_loop(self) -> None:
        """Maintain the MQTT connection with exponential backoff."""
        backoff = self._reconnect_interval
        while not self._stop_event.is_set():
            try:
                tls_context = self._build_tls_context() if self._tls_enabled else None
                async with Client(
                    hostname=self._broker,
                    port=self._port,
                    username=self._username,
                    password=self._password,
                    client_id=self._client_id,
                    keepalive=self._keepalive,
                    tls_context=tls_context,
                ) as active_client:
                    self._client = active_client
                    await self._apply_subscriptions(active_client)
                    self._connected_event.set()
                    self._last_message_at = time.monotonic()
                    backoff = self._reconnect_interval

                    async with active_client.unfiltered_messages() as messages:
                        await self._apply_subscriptions(active_client)
                        async for message in messages:
                            self._last_message_at = time.monotonic()
                            await self._dispatch_message(message)

            except MqttError as exc:
                logger.warning("MQTT connection lost: %s", exc)
                self._connected_event.clear()
                self._client = None
                if self._stop_event.is_set():
                    break
                await asyncio.sleep(backoff)
                backoff = min(backoff * 2, self._reconnect_interval_max)
            except Exception:  # pragma: no cover - defensive logging
                logger.exception("Unexpected error in MQTT connection loop.")
                self._connected_event.clear()
                self._client = None
                if self._stop_event.is_set():
                    break
                await asyncio.sleep(backoff)
                backoff = min(backoff * 2, self._reconnect_interval_max)
            else:
                # Connection closed gracefully; wait before reconnecting unless stopping.
                self._connected_event.clear()
                self._client = None
                if self._stop_event.is_set():
                    break
                await asyncio.sleep(backoff)
                backoff = min(backoff * 2, self._reconnect_interval_max)

    async def _dispatch_message(self, message: Message) -> None:
        """Dispatch incoming messages to registered handlers."""
        payload_bytes = message.payload
        try:
            payload = json.loads(payload_bytes.decode("utf-8"))
        except json.JSONDecodeError as exc:
            logger.warning("Discarding non-JSON payload on %s: %s", message.topic, exc)
            return

        try:
            self._validate_payload(payload, context=message.topic)
        except PayloadValidationError as exc:
            logger.warning("Invalid payload on %s: %s", message.topic, exc)
            return

        matching_handlers = [
            handler
            for subscription in self._subscriptions.values()
            if _topic_matches(subscription.topic, message.topic)
            for handler in subscription.handlers
        ]

        if not matching_handlers:
            logger.debug("No handler registered for topic %s", message.topic)
            return

        for handler in matching_handlers:
            try:
                result = handler(message.topic, payload)
                if inspect.isawaitable(result):
                    await result
            except Exception:  # pragma: no cover - defensive logging
                logger.exception("Handler error for topic %s", message.topic)

    async def _apply_subscriptions(self, client: Client) -> None:
        """Ensure all known subscriptions are active."""
        if not self._subscriptions:
            return
        subscriptions = [(sub.topic, sub.qos) for sub in self._subscriptions.values()]
        for topic, qos in subscriptions:
            await client.subscribe(topic, qos=qos)

    async def _watchdog_loop(self) -> None:
        """Periodically publish a heartbeat if traffic is silent."""
        await asyncio.sleep(self._watchdog_interval)
        while not self._stop_event.is_set():
            if self.is_connected:
                idle_seconds = time.monotonic() - self._last_message_at
                if idle_seconds >= self._watchdog_timeout:
                    heartbeat = {
                        "node_id": self._node_id,
                        "timestamp": _now_iso(),
                        "type": "watchdog",
                        "idle_seconds": round(idle_seconds, 2),
                    }
                    try:
                        await self.publish(
                            self._health_topic,
                            heartbeat,
                            qos=0,
                            retain=False,
                            ensure_base_fields=False,
                        )
                        self._last_message_at = time.monotonic()
                    except Exception:  # pragma: no cover - defensive logging
                        logger.exception("Failed to publish MQTT watchdog heartbeat.")
            await asyncio.sleep(self._watchdog_interval)

    def _build_tls_context(self) -> ssl.SSLContext | None:
        """Create an SSL context based on configuration."""
        ca_path = _resolve_optional_path(self._tls_ca)
        cert_path = _resolve_optional_path(self._tls_cert)
        key_path = _resolve_optional_path(self._tls_key)

        context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
        if ca_path:
            context.load_verify_locations(cafile=str(ca_path))
        if cert_path and key_path:
            context.load_cert_chain(certfile=str(cert_path), keyfile=str(key_path))
        if self._tls_insecure:
            context.check_hostname = False
            context.verify_mode = ssl.CERT_NONE

        return context

    def _validate_payload(
        self,
        payload: dict[str, Any],
        *,
        context: str,
        check_timestamp: bool = True,
    ) -> None:
        """Validate that a payload matches the expected schema."""
        error_prefix = f"[{context}] " if context else ""

        if not isinstance(payload, dict):
            raise PayloadValidationError(f"{error_prefix}Payload must be a JSON object.")

        missing = [field for field in self._required_fields if field not in payload]
        if missing:
            raise PayloadValidationError(f"{error_prefix}Missing fields {missing!r}")

        node_id = payload.get("node_id")
        if not isinstance(node_id, str) or not node_id.strip():
            raise PayloadValidationError(f"{error_prefix}node_id must be a non-empty string.")

        timestamp_raw = payload.get("timestamp")
        if not isinstance(timestamp_raw, str):
            raise PayloadValidationError(f"{error_prefix}timestamp must be an ISO-8601 string.")

        if check_timestamp:
            try:
                parsed = datetime.fromisoformat(timestamp_raw.replace("Z", "+00:00"))
            except ValueError as exc:  # pragma: no cover - defensive
                raise PayloadValidationError(
                    f"{error_prefix}timestamp must be an ISO-8601 string."
                ) from exc

            now = datetime.now(tz=UTC)
            if parsed > now + timedelta(seconds=30):
                raise PayloadValidationError(f"{error_prefix}timestamp is in the future.")


_default_client: RexMQTTClient | None = None


def get_mqtt_client() -> RexMQTTClient:
    """Return a shared RexMQTTClient instance."""
    global _default_client
    if _default_client is None:
        _default_client = RexMQTTClient()
    return _default_client