"""Module containing the `Client` class."""

import asyncio
import contextlib
import logging
import math
import typing
import uuid

if typing.TYPE_CHECKING:
    import datetime
    import types
    from collections.abc import AsyncIterator

    import google.protobuf.message

    from .types import Packet

import aiomqtt
import paho.mqtt.packettypes
import paho.mqtt.properties
import paho.mqtt.subscribeoptions

from .constants import (
    DEFAULT_MESSAGING_TIMEOUT,
    DEFAULT_MQTT_HOST,
    DEFAULT_MQTT_PORT,
    DEFAULT_RESPONSE_TOPIC_PREFIX,
)
from .exceptions import FastCCError, MessagingError, RouteExecutionError
from .qos import QoS
from .serialization import deserialize, serialize

_logger = logging.getLogger(__name__)


class Client:
    """Asynchronous MQTT client for FastCC.

    Although this client was specifically written for communication with
    a FastCC application, it can also be used without one. It provides
    general functions for subscribing to topics and publishing packages
    as a basic MQTT client would do. However, its true strength lies in
    the implemented request-response and stream mechanism which makes
    most sense when communicating with a FastCC application.

    Parameters
    ----------
    host
        MQTT broker host.
    port
        MQTT broker port.
    response_topic_prefix
        Base topic prefix to use for receiving responses.
    kwargs
        Additional keyword arguments to pass to `aiomqtt.Client`.
    """

    def __init__(
        self,
        host: str = DEFAULT_MQTT_HOST,
        port: int = DEFAULT_MQTT_PORT,
        *,
        response_topic_prefix: str = DEFAULT_RESPONSE_TOPIC_PREFIX,
        **kwargs: typing.Any,
    ) -> None:
        self._host = host
        self._port = port

        # Ensure MQTT v5 is used.
        kwargs.update({"protocol": aiomqtt.ProtocolVersion.V5})

        self._client = aiomqtt.Client(host, port, **kwargs)
        self._response_topic_prefix = response_topic_prefix
        self._dispatcher: asyncio.Task[None] | None = None
        self._messages: dict[str, asyncio.Queue[aiomqtt.Message]] = {}
        self._messages["default"] = asyncio.Queue()

    @property
    def host(self) -> str:
        """MQTT broker host."""
        return self._host

    @property
    def port(self) -> int:
        """MQTT broker port."""
        return self._port

    async def __aenter__(self) -> typing.Self:
        try:
            await self._client.__aenter__()
        except aiomqtt.MqttError:
            _logger.exception(
                "Failed to connect to MQTT broker on '%s:%d'",
                self._host,
                self._port,
            )
            error_message = "Failed to connect to MQTT broker"
            raise FastCCError(error_message) from None

        await self.__start_dispatcher()
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: types.TracebackType | None,
    ) -> None:
        await self.__stop_dispatcher()
        await self._client.__aexit__(exc_type, exc_value, traceback)

    @property
    def messages(self) -> AsyncIterator[aiomqtt.Message]:
        """Asynchronous iterator over incoming messages.

        This iterator excludes messages that are received on the
        specified response topic (and its subtopics) for the
        request-response mechanism. Those messages are handled
        internally and can be accessed through the `request` and
        `stream` methods.

        Warning
        -------
        Each access creates a new iterator consuming from the
        same internal queue. Calling this property multiple
        times will result in competing consumers where each
        message is delivered to only one iterator.

        Returns
        -------
        AsyncIterator[aiomqtt.Message]
            Asynchronous iterator over incoming messages.
        """

        async def _iter() -> AsyncIterator[aiomqtt.Message]:
            while True:
                yield await self._messages["default"].get()

        return _iter()

    async def publish(
        self,
        topic: str,
        packet: Packet | None = None,
        *,
        qos: QoS = QoS.AT_MOST_ONCE,
        retain: bool = False,
        properties: paho.mqtt.properties.Properties | None = None,
        timeout: datetime.timedelta | None = DEFAULT_MESSAGING_TIMEOUT,
    ) -> None:
        """Publish a packet to a topic.

        Parameters
        ----------
        topic
            Topic to publish the packet to.
        packet
            Packet to publish.
        qos
            QoS level to use for publishing the packet.
        retain
            Whether to publish the packet with the `retained` flag set
            to true.
        properties
            Properties to include with the publication.
        timeout
            Maximum time to wait for the publication. If `None`, wait
            indefinitely.

        Raises
        ------
        MessagingError
            If publishing to the topic failed or timed out.
        """
        payload = serialize(packet)
        timeout_seconds = _to_sec(timeout)

        try:
            # Use `timeout=math.inf` and rely on built-in timeout
            # handling using `asyncio.timeout`.
            async with asyncio.timeout(timeout_seconds):
                await self._client.publish(
                    topic,
                    payload,
                    qos=qos,
                    retain=retain,
                    properties=properties,
                    timeout=math.inf,
                )
        except aiomqtt.MqttCodeError as e:
            _logger.exception(
                "Publish to topic=%r with qos=%d (%s), retain=%r failed "
                "with error code: %r",
                topic,
                qos.value,
                qos.name,
                retain,
                e.rc,
            )
            error_message = (
                f"Publish to {topic!r} failed with error code: {e.rc!r}"
            )
            raise MessagingError(error_message) from e
        except TimeoutError as e:
            _logger.exception(
                "Publish to topic=%r with qos=%d (%s), retain=%r timed "
                "out after %.2f seconds",
                topic,
                qos.value,
                qos.name,
                retain,
                _to_float(timeout_seconds),
            )
            error_message = f"Publish to {topic!r} timed out"
            raise MessagingError(error_message) from e

        _logger.debug(
            "Published to topic=%r with qos=%d (%s), retain=%r: %r",
            topic,
            qos.value,
            qos.name,
            retain,
            payload,
        )

    async def subscribe(
        self,
        topic: str,
        *,
        qos: QoS = QoS.AT_MOST_ONCE,
        options: paho.mqtt.subscribeoptions.SubscribeOptions | None = None,
        properties: paho.mqtt.properties.Properties | None = None,
        timeout: datetime.timedelta | None = DEFAULT_MESSAGING_TIMEOUT,
    ) -> None:
        """Subscribe to a topic.

        Parameters
        ----------
        topic
            Topic to subscribe to.
        qos
            QoS level to use for subscribing.
        options
            Options to include with the subscription.
        properties
            Properties to include with the subscription.
        timeout
            Maximum time to wait for the subscription. If `None`, wait
            indefinitely.

        Raises
        ------
        MessagingError
            If subscribing to the topic failed or timed out.
        """
        timeout_seconds = _to_sec(timeout)

        if options is None:
            options = paho.mqtt.subscribeoptions.SubscribeOptions()

        # Parameter `qos` takes precedence over `options.qos` if both
        # are set.
        options.QoS = qos.value

        try:
            # Use `timeout=math.inf` and rely on built-in timeout
            # handling using `asyncio.timeout`.
            async with asyncio.timeout(timeout_seconds):
                await self._client.subscribe(
                    topic,
                    options=options,
                    properties=properties,
                    timeout=math.inf,
                )
        except aiomqtt.MqttCodeError as e:
            _logger.exception(
                "Subscribe to topic=%r with qos=%d (%s), failed with "
                "error code: %r",
                topic,
                qos.value,
                qos.name,
                e.rc,
            )
            error_message = (
                f"Subscribe to {topic!r} failed with error code: {e.rc!r}"
            )
            raise MessagingError(error_message) from e
        except TimeoutError as e:
            _logger.exception(
                "Subscribe to topic=%r with qos=%d (%s), timed out "
                "after %.2f seconds",
                topic,
                qos.value,
                qos.name,
                _to_float(timeout_seconds),
            )
            error_message = f"Subscribe to topic {topic!r} timed out"
            raise MessagingError(error_message) from e

        _logger.debug(
            "Subscribed to topic=%r with qos=%d (%s)",
            topic,
            qos.value,
            qos.name,
        )

    async def unsubscribe(
        self,
        topic: str,
        *,
        properties: paho.mqtt.properties.Properties | None = None,
        timeout: datetime.timedelta | None = DEFAULT_MESSAGING_TIMEOUT,
    ) -> None:
        """Unsubscribe from a topic.

        Parameters
        ----------
        topic
            Topic to unsubscribe from.
        properties
            Properties to include with the unsubscription.
        timeout
            Maximum time to wait for the unsubscription. If `None`, wait
            indefinitely.

        Raises
        ------
        MessagingError
            If unsubscribing to the topic failed or timed out.
        """
        timeout_seconds = _to_sec(timeout)

        try:
            # Use `timeout=math.inf` and rely on built-in timeout
            # handling using `asyncio.timeout`.
            async with asyncio.timeout(timeout_seconds):
                await self._client.unsubscribe(
                    topic,
                    properties=properties,
                    timeout=math.inf,
                )
        except aiomqtt.MqttCodeError as e:
            _logger.exception(
                "Unsubscribe from topic=%r failed with error code: %r",
                topic,
                e.rc,
            )
            error_message = (
                f"Unsubscribe from topic {topic!r} failed with error "
                f"code: {e.rc!r}"
            )
            raise MessagingError(error_message) from e
        except TimeoutError as e:
            _logger.exception(
                "Unsubscribe from topic=%r timed out after %.2f seconds",
                topic,
                _to_float(timeout_seconds),
            )
            error_message = f"Unsubscribe from {topic!r} timed out"
            raise MessagingError(error_message) from e

        _logger.debug("Unsubscribed from topic=%r", topic)

    @typing.overload
    async def request(
        self,
        topic: str,
        packet: Packet | None = None,
        *,
        response_type: None,
        qos: QoS = QoS.AT_LEAST_ONCE,
        retain: bool = False,
        properties: paho.mqtt.properties.Properties | None = None,
        timeout: datetime.timedelta | None = DEFAULT_MESSAGING_TIMEOUT,
    ) -> None: ...

    @typing.overload
    async def request(
        self,
        topic: str,
        packet: Packet | None = None,
        *,
        response_type: type[bytes],
        qos: QoS = QoS.AT_LEAST_ONCE,
        retain: bool = False,
        properties: paho.mqtt.properties.Properties | None = None,
        timeout: datetime.timedelta | None = DEFAULT_MESSAGING_TIMEOUT,
    ) -> bytes: ...

    @typing.overload
    async def request(
        self,
        topic: str,
        packet: Packet | None = None,
        *,
        response_type: type[str],
        qos: QoS = QoS.AT_LEAST_ONCE,
        retain: bool = False,
        properties: paho.mqtt.properties.Properties | None = None,
        timeout: datetime.timedelta | None = DEFAULT_MESSAGING_TIMEOUT,
    ) -> str: ...

    @typing.overload
    async def request(
        self,
        topic: str,
        packet: Packet | None = None,
        *,
        response_type: type[int],
        qos: QoS = QoS.AT_LEAST_ONCE,
        retain: bool = False,
        properties: paho.mqtt.properties.Properties | None = None,
        timeout: datetime.timedelta | None = DEFAULT_MESSAGING_TIMEOUT,
    ) -> int: ...

    @typing.overload
    async def request(
        self,
        topic: str,
        packet: Packet | None = None,
        *,
        response_type: type[float],
        qos: QoS = QoS.AT_LEAST_ONCE,
        retain: bool = False,
        properties: paho.mqtt.properties.Properties | None = None,
        timeout: datetime.timedelta | None = DEFAULT_MESSAGING_TIMEOUT,
    ) -> float: ...

    @typing.overload
    async def request[T: google.protobuf.message.Message](
        self,
        topic: str,
        packet: Packet | None = None,
        *,
        response_type: type[T],
        qos: QoS = QoS.AT_LEAST_ONCE,
        retain: bool = False,
        properties: paho.mqtt.properties.Properties | None = None,
        timeout: datetime.timedelta | None = DEFAULT_MESSAGING_TIMEOUT,
    ) -> T: ...

    async def request(  # type: ignore[no-untyped-def]
        self,
        topic: str,
        packet: Packet | None = None,
        *,
        response_type: typing.Any = None,
        qos: QoS = QoS.AT_LEAST_ONCE,
        retain: bool = False,
        properties: paho.mqtt.properties.Properties | None = None,
        timeout: datetime.timedelta | None = DEFAULT_MESSAGING_TIMEOUT,
    ):
        """Publish a packet and wait for a single response.

        Implements the request-response pattern using MQTT v5 response
        topics. A unique response topic is generated, subscribed to,
        and set as a property on the published message. The server is
        expected to publish exactly one response to that topic.

        Parameters
        ----------
        topic
            Topic to publish the request to.
        packet
            Packet to publish as the request payload.
        response_type
            Expected type of the response packet for deserialization.
            If `None`, the response payload is expected to be empty.
        qos
            QoS level to use for the request and response.
        retain
            Whether to publish the request with the `retained` flag set
            to true.
        properties
            Properties to include with the request publication.
        timeout
            Maximum time to wait for the response. If `None`, wait
            indefinitely.

        Returns
        -------
        T | None
            Deserialized response packet, or `None` if the server
            responded with an empty payload.

        Raises
        ------
        MessagingError
            If the request failed or timed out.
        """
        timeout_seconds = _to_sec(timeout)

        properties = _validate_properties(properties, qos, context="request")

        correlation_id = uuid.uuid4().hex
        properties.ResponseTopic = (
            f"{self._response_topic_prefix}/{correlation_id}"
        )
        properties.CorrelationData = correlation_id.encode()

        queue: asyncio.Queue[aiomqtt.Message] = asyncio.Queue()
        self._messages[correlation_id] = queue
        try:
            await self.publish(
                topic,
                packet,
                qos=qos,
                retain=retain,
                properties=properties,
                timeout=timeout,
            )

            try:
                msg = await asyncio.wait_for(
                    queue.get(),
                    timeout=timeout_seconds,
                )
            except TimeoutError as e:
                _logger.exception(
                    "Request to topic=%r with qos=%d (%s), retain=%r "
                    "timed out after waiting %.2f seconds for response",
                    topic,
                    qos.value,
                    qos.name,
                    retain,
                    _to_float(timeout_seconds),
                )
                error_message = (
                    f"Request to topic {topic!r} timed out waiting for response"
                )
                raise MessagingError(error_message) from e

            _check_error_response(msg, topic=topic, context="request")

            if response_type is None:
                if msg.payload:
                    _logger.warning(
                        "Received non-empty response for request to "
                        "topic=%r with response_type=None: %r",
                        topic,
                        msg.payload,
                    )
                return None

            response = deserialize(msg.payload, response_type)
            _logger.debug(
                "Received response for request to topic=%r: %r",
                topic,
                response,
            )
            return response
        finally:
            self._messages.pop(correlation_id, None)

    @typing.overload
    def stream(
        self,
        topic: str,
        packet: Packet | None = ...,
        *,
        response_type: type[bytes] = ...,
        qos: QoS = ...,
        retain: bool = ...,
        properties: paho.mqtt.properties.Properties | None = ...,
        timeout: datetime.timedelta | None = ...,
    ) -> AsyncIterator[bytes]: ...

    @typing.overload
    def stream(
        self,
        topic: str,
        packet: Packet | None = ...,
        *,
        response_type: type[str] = ...,
        qos: QoS = ...,
        retain: bool = ...,
        properties: paho.mqtt.properties.Properties | None = ...,
        timeout: datetime.timedelta | None = ...,
    ) -> AsyncIterator[str]: ...

    @typing.overload
    def stream(
        self,
        topic: str,
        packet: Packet | None = ...,
        *,
        response_type: type[int] = ...,
        qos: QoS = ...,
        retain: bool = ...,
        properties: paho.mqtt.properties.Properties | None = ...,
        timeout: datetime.timedelta | None = ...,
    ) -> AsyncIterator[int]: ...

    @typing.overload
    def stream(
        self,
        topic: str,
        packet: Packet | None = ...,
        *,
        response_type: type[float] = ...,
        qos: QoS = ...,
        retain: bool = ...,
        properties: paho.mqtt.properties.Properties | None = ...,
        timeout: datetime.timedelta | None = ...,
    ) -> AsyncIterator[float]: ...

    @typing.overload
    def stream[T: google.protobuf.message.Message](
        self,
        topic: str,
        packet: Packet | None = ...,
        *,
        response_type: type[T] = ...,
        qos: QoS = ...,
        retain: bool = ...,
        properties: paho.mqtt.properties.Properties | None = ...,
        timeout: datetime.timedelta | None = ...,
    ) -> AsyncIterator[T]: ...

    async def stream(  # type: ignore[no-untyped-def]
        self,
        topic,
        packet=None,
        *,
        response_type,
        qos=QoS.AT_LEAST_ONCE,
        retain=False,
        properties=None,
        timeout=DEFAULT_MESSAGING_TIMEOUT,
    ):
        """Publish a packet and stream multiple responses.

        Implements the streaming pattern using MQTT v5 response topics.
        A unique response topic is generated, subscribed to, and set as
        a property on the published message. The server may publish
        multiple responses to that topic. The stream ends when a
        response with an empty payload (`b""`) is received.

        Parameters
        ----------
        topic
            Topic to publish the stream request to.
        packet
            Packet to publish as the request payload.
        response_type
            Expected type of the response packets for deserialization.
        qos
            QoS level to use for the request and responses.
        retain
            Whether to publish the request with the `retained` flag
            set to true.
        properties
            Properties to include with the request publication.
        timeout
            Maximum time to wait for each individual response. If
            `None`, wait indefinitely.

        Yields
        ------
        T
            Deserialized response packets as they arrive.

        Raises
        ------
        MessagingError
            If the stream request failed or a response timed out.
        """
        timeout_seconds = _to_sec(timeout)
        properties = _validate_properties(properties, qos, context="stream")

        correlation_id = uuid.uuid4().hex
        properties.ResponseTopic = (
            f"{self._response_topic_prefix}/{correlation_id}"
        )
        properties.CorrelationData = correlation_id.encode()

        queue: asyncio.Queue[aiomqtt.Message] = asyncio.Queue()
        self._messages[correlation_id] = queue
        try:
            await self.publish(
                topic,
                packet,
                qos=qos,
                retain=retain,
                properties=properties,
                timeout=timeout,
            )

            while True:
                try:
                    msg = await asyncio.wait_for(
                        queue.get(),
                        timeout=timeout_seconds,
                    )
                except TimeoutError as e:
                    _logger.exception(
                        "Stream from topic=%r with qos=%d (%s), "
                        "retain=%r timed out after waiting %.2f "
                        "seconds for next response",
                        topic,
                        qos.value,
                        qos.name,
                        retain,
                        _to_float(timeout_seconds),
                    )
                    error_message = (
                        f"Stream from topic {topic!r} timed out "
                        f"waiting for next response"
                    )
                    raise MessagingError(error_message) from e

                _check_error_response(msg, topic=topic, context="stream")

                if not msg.payload:
                    _logger.debug("Stream from topic=%r ended", topic)
                    return

                response = deserialize(msg.payload, response_type)
                _logger.debug(
                    "Received stream response for topic=%r: %r",
                    topic,
                    response,
                )
                yield response
        finally:
            self._messages.pop(correlation_id, None)

    async def __start_dispatcher(self) -> None:
        await self.subscribe(
            f"{self._response_topic_prefix}/#",
            qos=QoS.AT_LEAST_ONCE,
            options=paho.mqtt.subscribeoptions.SubscribeOptions(
                noLocal=True,
            ),
        )

        self._dispatcher = asyncio.create_task(self.__dispatch())

    async def __stop_dispatcher(self) -> None:
        if self._dispatcher is not None:
            self._dispatcher.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await self._dispatcher
            self._dispatcher = None

    async def __dispatch(self) -> None:
        try:
            async for msg in self._client.messages:
                topic = msg.topic.value
                if topic.startswith(self._response_topic_prefix):
                    correlation_id = topic.rsplit("/", 1)[-1]
                    queue = self._messages.get(correlation_id)
                    if queue is None:
                        _logger.warning(
                            "Received message with unknown correlation ID: %r",
                            correlation_id,
                        )
                        continue
                    await queue.put(msg)
                else:
                    await self._messages["default"].put(msg)

        except asyncio.CancelledError:
            raise
        except Exception:
            _logger.exception("Dispatcher encountered an unexpected error")
            raise


def _to_float(timeout: float | None) -> float:
    """Convert a timeout value to a float, using `math.inf` for `None`.

    Parameters
    ----------
    timeout
        Timeout value to convert, or `None` for no timeout.

    Returns
    -------
    float
        The timeout value as a float, or `math.inf` if `timeout` is `None`.
    """
    return timeout if timeout is not None else math.inf


def _to_sec(timeout: datetime.timedelta | None) -> float | None:
    """Convert a timeout value to seconds, or `None` for no timeout.

    Parameters
    ----------
    timeout
        Timeout value to convert, or `None` for no timeout.

    Returns
    -------
    float | None
        The timeout value in seconds, or `None` if `timeout` is `None`.
    """
    return timeout.total_seconds() if timeout is not None else None


def _validate_properties(
    properties: paho.mqtt.properties.Properties | None,
    qos: QoS,
    context: str,
) -> paho.mqtt.properties.Properties:
    """Validate properties for request and stream calls.

    Parameters
    ----------
    properties
        Properties to validate, or `None` to create default properties.
    qos
        Quality of Service level for the call.
    context
        Description for log messages.

    Returns
    -------
    paho.mqtt.properties.Properties
        Validated properties to use for the call.

    Raises
    ------
    MessagingError
        If the properties are invalid for the call.
    """
    if qos == QoS.AT_MOST_ONCE:
        _logger.error(
            "QoS level %r is not supported for %r calls, must be at least %r",
            qos.name,
            context,
            QoS.AT_LEAST_ONCE.name,
        )
        error_message = f"Invalid QoS level for {context!r}: {qos.name}."
        raise MessagingError(error_message)

    if properties is None:
        properties = paho.mqtt.properties.Properties(
            paho.mqtt.packettypes.PacketTypes.PUBLISH,
        )

    if properties.packetType != paho.mqtt.packettypes.PacketTypes.PUBLISH:
        name = paho.mqtt.packettypes.PacketTypes.Names[properties.packetType]
        _logger.error(
            "Invalid packet type for %r calls: %s (must be %r)",
            context,
            name,
            paho.mqtt.packettypes.PacketTypes.Names[
                paho.mqtt.packettypes.PacketTypes.PUBLISH
            ],
        )
        error_message = f"Invalid PacketType: {name}"
        raise MessagingError(error_message)

    if hasattr(properties, "ResponseTopic"):
        _logger.error(
            "ResponseTopic property must not be set for %r calls",
            context,
        )
        error_message = "ResponseTopic must not be set"
        raise MessagingError(error_message)

    if hasattr(properties, "CorrelationData"):
        _logger.error(
            "CorrelationData property must not be set for %r calls",
            context,
        )
        error_message = "CorrelationData must not be set"
        raise MessagingError(error_message)

    return properties


def _check_error_response(
    message: aiomqtt.Message,
    *,
    topic: str,
    context: str,
) -> None:
    """Raise if the response message contains an error.

    Parameters
    ----------
    message
        The response message to check.
    topic
        The original request/stream topic, used in error and log
        messages.
    context
        Description for log messages.

    Raises
    ------
    RouteExecutionError
        If the message contains an error user property.
    """
    user_properties = getattr(message.properties, "UserProperty", [])
    raw_error_code = next(
        (t[1] for t in user_properties if t[0] == "error"),
        None,
    )
    if raw_error_code is None:
        return

    error_code = int(raw_error_code)
    error_message = deserialize(message.payload, str)
    details = "Received error response [error_code=%r] for %s on topic %r: %s"
    _logger.error(
        details,
        error_code,
        context,
        topic,
        error_message,
    )
    raise RouteExecutionError(
        error_message,
        error_code=error_code,
    )
