"""
WebSocket Connection Pool

Manages WebSocket connections with:
- Automatic reconnection (exponential backoff with jitter)
- Active subscription restoration after reconnect
- Proper cleanup on close
"""

from __future__ import annotations

import asyncio
import random
from dataclasses import dataclass, field
from typing import Callable

import websockets
import websockets.client

from dvgateway.auth.manager import AuthManager
from dvgateway.types import Logger, ReconnectOptions


@dataclass
class WsSubscription:
    id: str
    path: str
    params: dict[str, str] = field(default_factory=dict)
    on_message: Callable[[bytes], None] = lambda _: None
    on_close: Callable[[], None] | None = None


class WsPool:
    """WebSocket connection pool with automatic reconnection."""

    def __init__(self, auth: AuthManager, reconnect: ReconnectOptions, logger: Logger) -> None:
        self._auth = auth
        self._logger = logger
        self._reconnect = reconnect

        self._connections: dict[str, _ManagedConnection] = {}
        self._event_handlers: dict[str, list[Callable[..., None]]] = {}

    def on(self, event: str, handler: Callable[..., None]) -> None:
        """Register an event handler."""
        self._event_handlers.setdefault(event, []).append(handler)

    def once(self, event: str, handler: Callable[..., None]) -> None:
        """Register a one-time event handler."""
        def wrapper(*args: object) -> None:
            handler(*args)
            handlers = self._event_handlers.get(event, [])
            if wrapper in handlers:
                handlers.remove(wrapper)

        self.on(event, wrapper)

    def emit(self, event: str, *args: object) -> None:
        """Emit an event to all registered handlers."""
        for handler in self._event_handlers.get(event, []):
            handler(*args)

    async def subscribe(self, sub: WsSubscription) -> None:
        """Open a managed WebSocket connection for a subscription."""
        conn = _ManagedConnection(subscription=sub)
        self._connections[sub.id] = conn
        await self._connect(conn)

    def unsubscribe(self, sub_id: str) -> None:
        """Close a specific subscription."""
        conn = self._connections.get(sub_id)
        if not conn:
            return
        conn.closed = True
        if conn.reconnect_task and not conn.reconnect_task.done():
            conn.reconnect_task.cancel()
        if conn.ws:
            asyncio.ensure_future(conn.ws.close(1000, "unsubscribed"))
        del self._connections[sub_id]
        self._logger.debug({"id": sub_id}, "WebSocket subscription closed")

    def close_all(self) -> None:
        """Close all connections."""
        for sub_id in list(self._connections.keys()):
            self.unsubscribe(sub_id)

    async def send(self, sub_id: str, data: bytes) -> None:
        """Send data through a subscription's WebSocket."""
        conn = self._connections.get(sub_id)
        if conn and conn.ws:
            await conn.ws.send(data)

    async def _connect(self, conn: _ManagedConnection) -> None:
        if conn.closed:
            return

        try:
            url = await self._auth.build_ws_url(conn.subscription.path, conn.subscription.params)
        except Exception as e:
            self._logger.error({"err": str(e), "id": conn.subscription.id}, "Failed to build WebSocket URL")
            self._schedule_reconnect(conn)
            return

        self._logger.debug({"id": conn.subscription.id, "path": conn.subscription.path}, "Opening WebSocket")

        try:
            ws = await websockets.client.connect(url, open_timeout=10)
        except Exception as e:
            self._logger.error({"err": str(e), "id": conn.subscription.id}, "WebSocket connection failed")
            self._schedule_reconnect(conn)
            return

        conn.ws = ws
        conn.reconnect_attempts = 0
        self._logger.debug({"id": conn.subscription.id}, "WebSocket connected")
        self.emit("connected", conn.subscription.id)

        # Start reading messages in the background
        conn.read_task = asyncio.ensure_future(self._read_loop(conn))

    async def _read_loop(self, conn: _ManagedConnection) -> None:
        assert conn.ws is not None
        try:
            async for message in conn.ws:
                if conn.closed:
                    break
                if isinstance(message, bytes):
                    conn.subscription.on_message(message)
                else:
                    conn.subscription.on_message(message.encode("utf-8"))
        except websockets.exceptions.ConnectionClosed:
            pass
        except Exception as e:
            self._logger.error({"id": conn.subscription.id, "err": str(e)}, "WebSocket error")
        finally:
            if not conn.closed:
                self._logger.warn({"id": conn.subscription.id}, "WebSocket closed unexpectedly")
                if conn.subscription.on_close:
                    conn.subscription.on_close()
                self._schedule_reconnect(conn)

    def _schedule_reconnect(self, conn: _ManagedConnection) -> None:
        if conn.closed:
            return
        max_attempts = self._reconnect.max_attempts
        if conn.reconnect_attempts >= max_attempts:
            self._logger.error(
                {"id": conn.subscription.id, "attempts": conn.reconnect_attempts},
                "Max reconnect attempts reached",
            )
            if self._reconnect.on_max_attempts_reached:
                self._reconnect.on_max_attempts_reached()
            self.emit("maxAttemptsReached", conn.subscription.id)
            return

        jitter = random.random() * 0.3
        base = self._reconnect.initial_delay_ms * (self._reconnect.backoff_multiplier ** conn.reconnect_attempts)
        delay_ms = min(base * (1 + jitter), self._reconnect.max_delay_ms)
        conn.reconnect_attempts += 1

        self._logger.info(
            {"id": conn.subscription.id, "attempt": conn.reconnect_attempts, "delay_ms": round(delay_ms)},
            "Scheduling WebSocket reconnect",
        )

        async def _reconnect_after_delay() -> None:
            await asyncio.sleep(delay_ms / 1000.0)
            if self._reconnect.on_reconnect:
                self._reconnect.on_reconnect(conn.reconnect_attempts)
            self.emit("reconnecting", conn.subscription.id, conn.reconnect_attempts)
            await self._connect(conn)

        conn.reconnect_task = asyncio.ensure_future(_reconnect_after_delay())


@dataclass
class _ManagedConnection:
    subscription: WsSubscription = field(default_factory=WsSubscription)
    ws: websockets.client.WebSocketClientProtocol | None = None  # type: ignore[type-arg]
    reconnect_attempts: int = 0
    closed: bool = False
    reconnect_task: asyncio.Task[None] | None = None
    read_task: asyncio.Task[None] | None = None
