from __future__ import annotations

import asyncio
import json
import random
from collections.abc import Awaitable, Callable
from typing import Any

import aiohttp

from yuxi.utils.logging_config import logger


_MAX_ERROR_MSG_LEN = 200


class IMessageMonitor:
    def __init__(self, config: dict[str, Any], bridge):
        self._config = config
        self._bridge = bridge
        self._message_handler: Callable[[dict[str, Any]], Awaitable[None]] | None = None
        self._task: asyncio.Task | None = None
        self._reconnect_delay = 5.0
        self._initial_reconnect_delay = config.get("ws_reconnect_initial_delay", 5.0)
        self._max_reconnect_delay = config.get("ws_reconnect_max_delay", 60.0)
        self._max_subscribe_attempts = config.get(
            "watchSubscribeMaxAttempts", config.get("watch_subscribe_max_attempts", 3)
        )
        self._subscribe_retry_delay_ms = config.get(
            "watchSubscribeRetryDelayMs", config.get("watch_subscribe_retry_delay_ms", 1000)
        )
        self._subscribe_attempts = 0

    def on_message(self, handler: Callable[[dict[str, Any]], Awaitable[None]]) -> None:
        self._message_handler = handler

    async def start(self) -> None:
        if self._task and not self._task.done():
            return
        self._task = asyncio.create_task(self._ws_loop())

    async def stop(self) -> None:
        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
            self._task = None

    async def _ws_loop(self) -> None:
        while True:
            try:
                self._subscribe_attempts = 0
                await self._connect_ws()
            except asyncio.CancelledError:
                break
            except Exception as e:
                self._subscribe_attempts += 1
                if self._subscribe_attempts >= self._max_subscribe_attempts:
                    logger.error(
                        f"[iMessage] Max subscribe attempts ({self._max_subscribe_attempts}) reached, stopping monitor"
                    )
                    break
                delay = self._reconnect_delay * (0.5 + random.random())
                err_msg = str(e)[:_MAX_ERROR_MSG_LEN]
                logger.error(
                    f"[iMessage] WebSocket error (attempt {self._subscribe_attempts}), "
                    f"reconnecting in {delay:.1f}s: {err_msg}"
                )
                await asyncio.sleep(delay)
                self._reconnect_delay = min(self._reconnect_delay * 1.5, self._max_reconnect_delay)

    async def _connect_ws(self) -> None:
        url = self._bridge.ws_url
        logger.info(f"[iMessage] Connecting to BlueBubbles WebSocket: {url}")

        async with aiohttp.ClientSession() as session:
            headers = {"X-BlueBubbles-Password": self._bridge.password}
            async with session.ws_connect(url, headers=headers) as ws:
                self._reconnect_delay = self._initial_reconnect_delay
                logger.info("[iMessage] BlueBubbles WebSocket connected")

                async for msg in ws:
                    if msg.type == aiohttp.WSMsgType.TEXT:
                        try:
                            data = json.loads(msg.data)
                            event_type = data.get("type", "")
                            logger.debug(f"[iMessage] WS event: {event_type}")

                            if event_type == "error":
                                err_data = data.get("data", {})
                                err_str = str(err_data)
                                if isinstance(err_data, dict):
                                    err_str = err_data.get("message", err_data.get("error", str(err_data)))
                                logger.error(f"[iMessage] BlueBubbles error: {err_str[:_MAX_ERROR_MSG_LEN]}")
                            elif self._message_handler:
                                await self._message_handler(data)
                        except json.JSONDecodeError:
                            logger.warning(f"[iMessage] WS non-JSON message: {msg.data[:200]}")
                    elif msg.type == aiohttp.WSMsgType.ERROR:
                        logger.error(f"[iMessage] WS error: {ws.exception()}")
                        break
                    elif msg.type == aiohttp.WSMsgType.CLOSED:
                        logger.info("[iMessage] BlueBubbles WebSocket closed")
                        break