"""AWS IoT MQTT implementation"""

import asyncio
import contextlib
import functools
import itertools
import json
import logging
import re
import threading
import uuid
from collections.abc import Callable

from awscrt.auth import AwsSigningConfig, aws_sign_request
from awscrt.http import HttpHeaders, HttpRequest
from paho.mqtt.client import CallbackAPIVersion, Client

from ._api import PYUPGW_AWS_IOT_ENDPOINT, PYUPGW_AWS_REGION, AwsApi
from .errors import ClientError

logger = logging.getLogger(__name__)

AWS_IOTDEVICEGATEWAY_SERVICE = "iotdevicegateway"
AWS_IOTDEVICEGATEWAY_MQTT_PATH = "/mqtt"
AWS_TOPIC_GET = "get"
AWS_TOPIC_UPDATE = "update"
AWS_TOPIC_ACCEPTED = "accepted"
AWS_TOPIC_REJECTED = "rejected"
AWS_TOPIC_RE = re.compile(
    r"\$aws/things/([\w-]+)/shadow/(get|update)/(accepted|rejected)"
)
MQTT_KEEPALIVE = 30
PUBLISHING_TIMEOUT = 60


def _aws_shadow_topic(thing_name: str, command: str, result: str | None = None):
    suffix = f"/{result}" if result is not None else ""
    return f"$aws/things/{thing_name}/shadow/{command}{suffix}"


class IotShadowMqtt(contextlib.AbstractAsyncContextManager):
    """MQTT client that interacts with the AWS IoT shadow service

    Spawns a Paho MQTT based network loop that will:
    - Handle connecting to AWS IoT via MQTT over WebSocket
    - Handle sourcing credentials
    - Handle reconnecting, and respawning connection from time to time
    - Give asyncio interface to interacting with the network thread
    """

    def __init__(  # noqa: PLR0913
        self,
        *,
        aws: AwsApi,
        identity_id: str,
        client_name: str,
        thing_names: list[str],
        loop: asyncio.AbstractEventLoop,
        on_response_state_received: Callable[[str, dict | None], None],
        on_disconnected: Callable[[list[str]], None],
        on_connected: Callable[[], None],
    ):
        self._aws = aws
        self._identity_id = identity_id
        self._client_name = client_name
        self._thing_names = thing_names
        self._on_response_state_received = on_response_state_received
        self._on_disconnected = on_disconnected
        self._on_connected = on_connected
        self._loop = loop
        self._client: Client | None = None
        self._thread = threading.Thread(target=self._mqtt_loop)
        self._publish_lock = threading.Lock()
        self._subscription_lock = threading.Lock()
        self._async_initialized_event = asyncio.Event()
        self._async_quit_done_event = asyncio.Event()
        self._pending_subscriptions = 0
        self._pending_publish_futures: dict[str, asyncio.Future] = {}
        self._pending_tasks: set[asyncio.Task] = set()

    async def __aenter__(self):
        self._thread.start()
        await self._async_initialized_event.wait()
        return self

    async def __aexit__(self, exc_type, exc_value, traceback):
        for task in self._pending_tasks:
            task.cancel()
        for future in self._pending_publish_futures.values():
            future.cancel()
        self._client.disconnect()
        await self._async_quit_done_event.wait()
        self._thread.join()

    async def get(self, thing_name: str):
        """Publish message to get topic for ``thing_name`` and await response"""
        return await self._publish(thing_name, "get", {})

    async def update(self, thing_name: str, payload: dict):
        """Publish ``payload`` to update topic for ``thing_name`` and await response"""
        return await self._publish(thing_name, "update", payload)

    def _build_headers(self, headers):
        self._aws.check_token()
        credentials_provider = self._aws.get_credentials_provider(self._identity_id)
        signing_config = AwsSigningConfig(
            credentials_provider=credentials_provider,
            region=PYUPGW_AWS_REGION,
            service=AWS_IOTDEVICEGATEWAY_SERVICE,
        )
        http_request = HttpRequest(
            path=AWS_IOTDEVICEGATEWAY_MQTT_PATH,
            headers=HttpHeaders(list(headers.items())),
        )
        signed_http_request = aws_sign_request(
            http_request=http_request,
            signing_config=signing_config,
        ).result()
        return dict(signed_http_request.headers)

    def _mqtt_loop(self):
        client = Client(
            CallbackAPIVersion.VERSION2,
            client_id=f"{self._client_name}-{uuid.uuid4()}",
            transport="websockets",
        )
        client.ws_set_options(
            path=AWS_IOTDEVICEGATEWAY_MQTT_PATH, headers=self._build_headers
        )
        client.tls_set()
        client.on_message = self._on_message
        client.on_connect = self._on_connect
        client.on_subscribe = self._on_subscribe
        client.on_disconnect = self._on_disconnect
        client.connect(PYUPGW_AWS_IOT_ENDPOINT, 443, keepalive=MQTT_KEEPALIVE)
        client.loop_forever()
        self._loop.call_soon_threadsafe(self._async_quit_done_event.set)

    def _on_connect(self, client, _userdata, _flags, reason_code, _properties):
        if reason_code.is_failure:
            logger.warning("MQTT client failed to connect with code %r", reason_code)
            return
        logger.info("MQTT client connected")
        with self._subscription_lock:
            self._pending_subscriptions = len(self._thing_names)
            logger.debug("Subscribing for %d devices", self._pending_subscriptions)
            for thing_name in self._thing_names:
                client.subscribe(
                    [
                        (_aws_shadow_topic(thing_name, command, result), 0)
                        for (command, result) in itertools.product(
                            [AWS_TOPIC_GET, AWS_TOPIC_UPDATE],
                            [AWS_TOPIC_ACCEPTED, AWS_TOPIC_REJECTED],
                        )
                    ]
                )

    def _on_subscribe(self, client, _userdata, _mid, reason_codes, _properties):
        if any(reason_code.is_failure for reason_code in reason_codes):
            logger.warning(
                "MQTT client failed to subscribe with codes %r", reason_codes
            )
            return
        with self._subscription_lock:
            self._pending_subscriptions -= 1
            logger.debug(
                "MQTT client subscription completed, pending %d subscriptions",
                self._pending_subscriptions,
            )
            if self._pending_subscriptions > 0:
                return
        with self._publish_lock:
            self._client = client
        # We are here after reconnecting. Request new state from devices.
        if self._async_initialized_event.is_set():
            for thing_name in self._thing_names:
                self._loop.call_soon_threadsafe(self._async_get, thing_name)
        self._loop.call_soon_threadsafe(self._async_initialized_event.set)
        self._loop.call_soon_threadsafe(self._on_connected)

    def _on_message(self, _client, _userdata, message):
        logger.debug("Received message from %s: %r", message.topic, message.payload)
        parsed_topic = AWS_TOPIC_RE.fullmatch(message.topic)
        if parsed_topic:
            thing_name, _command, result = parsed_topic.groups()
            try:
                parsed_payload = json.loads(message.payload)
            except json.JSONDecodeError:
                logger.warning("Unable to parse MQTT message: %r", message.payload)
                parsed_payload = None
            if result == AWS_TOPIC_ACCEPTED:
                self._loop.call_soon_threadsafe(
                    self._accepted_callback, thing_name, parsed_payload
                )
            elif result == AWS_TOPIC_REJECTED:
                self._loop.call_soon_threadsafe(self._rejected_callback, parsed_payload)

    def _on_disconnect(self, _client, _userdata, _flags, reason_code, _properties):
        if reason_code.is_failure:
            logger.warning(
                "MQTT client unexpectedly disconnected with code %r", reason_code
            )
        else:
            logger.info("MQTT client disconnected")
        self._loop.call_soon_threadsafe(self._on_disconnected, self._thing_names)

    def _accepted_callback(self, thing_name: str, payload: dict | None):
        if publish_future := self._get_pending_publish_future_for_response(payload):
            publish_future.set_result(payload)
        self._on_response_state_received(thing_name, payload)

    def _rejected_callback(self, payload: dict | None):
        if publish_future := self._get_pending_publish_future_for_response(payload):
            publish_future.set_exception(ClientError(f"Message rejected: {payload!r}"))

    def _get_pending_publish_future_for_response(self, payload: dict | None):
        client_token = payload.pop("clientToken", None) if payload else None
        if (
            client_token
            and (publish_future := self._pending_publish_futures.get(client_token))
            and not publish_future.done()
        ):
            return publish_future
        return None

    async def _publish(self, thing_name: str, command: str, payload: dict):
        topic = _aws_shadow_topic(thing_name, command)
        client_token = str(uuid.uuid4())
        payload["clientToken"] = client_token
        publish_future = self._loop.create_future()
        self._pending_publish_futures[client_token] = publish_future
        try:
            async with asyncio.timeout(PUBLISHING_TIMEOUT):
                await asyncio.to_thread(self._do_publish, topic, json.dumps(payload))
                return await publish_future
        finally:
            self._pending_publish_futures.pop(client_token, None)

    def _do_publish(self, topic: str, payload: str):
        with self._publish_lock:
            if not self._client:
                raise ClientError("MQTT client not initialized")
            logger.debug("Publishing message to %s: %r", topic, payload)
            self._client.publish(topic, payload)

    def _async_get(self, thing_name: str):
        task = self._loop.create_task(self.get(thing_name))
        self._pending_tasks.add(task)
        task.add_done_callback(functools.partial(self._async_get_done, thing_name))

    def _async_get_done(self, thing_name: str, task: asyncio.Task):
        if exc := task.exception():
            logger.warning("Unable to get state for %s", thing_name, exc_info=exc)
        self._pending_tasks.discard(task)
