"""Async WebSocket client for real-time Sibilance events.

Connects to the lounge server WebSocket and dispatches typed events
to a callback. Handles reconnection with exponential backoff.
"""

from __future__ import annotations

import asyncio
import json
import logging
from enum import Enum
from typing import Any, Callable, Protocol
from urllib.parse import urlparse, urlunparse

import websockets
import websockets.asyncio.client as ws_client

from sibilance.auth.credentials import read_credentials

log = logging.getLogger(__name__)


class WSEventType(str, Enum):
    CONNECTED = "connected"
    AGENT_STATE_UPDATE = "agent_state_update"
    LOCATION_UPDATE = "location_update"
    CONVERSATION_UPDATE = "conversation_update"
    JOURNAL_ENTRY = "journal_entry"
    WORLD_STATE_UPDATE = "world_state_update"
    MYSTERY_CLUE = "mystery_clue"
    RARE_EVENT = "rare_event"
    KAFKA_EVENT = "event"
    PONG = "pong"


class WSEvent:
    """A parsed WebSocket event from the server."""

    __slots__ = ("type", "data")

    def __init__(self, event_type: WSEventType, data: dict[str, Any] | None = None):
        self.type = event_type
        self.data = data or {}


class WSEventHandler(Protocol):
    """Callback protocol for WebSocket events."""

    def __call__(self, event: WSEvent) -> None: ...


class SibilanceWebSocket:
    """Manages a WebSocket connection to the Sibilance server.

    Usage:
        ws = SibilanceWebSocket(server_url, agent_id, on_event=handler)
        task = asyncio.create_task(ws.run())
        # ... later ...
        await ws.close()
    """

    def __init__(
        self,
        server_url: str,
        agent_id: str,
        *,
        on_event: WSEventHandler | None = None,
        on_connected: Callable[[], None] | None = None,
        on_disconnected: Callable[[], None] | None = None,
    ) -> None:
        self._server_url = server_url
        self._agent_id = agent_id
        self._on_event = on_event
        self._on_connected = on_connected
        self._on_disconnected = on_disconnected
        self._ws: ws_client.ClientConnection | None = None
        self._running = False
        self._reconnect_delay = 1.0
        self._max_reconnect_delay = 30.0

    @property
    def is_connected(self) -> bool:
        return self._ws is not None

    def _ws_url(self) -> str:
        """Convert HTTP(S) URL to WS(S) URL."""
        parsed = urlparse(self._server_url)
        scheme = "wss" if parsed.scheme == "https" else "ws"
        return urlunparse(parsed._replace(scheme=scheme))

    async def run(self) -> None:
        """Connect and process messages until close() is called.

        Reconnects automatically with exponential backoff on disconnection.
        """
        self._running = True
        while self._running:
            try:
                await self._connect_and_listen()
            except (
                websockets.exceptions.ConnectionClosed,
                websockets.exceptions.WebSocketException,
                OSError,
            ) as exc:
                log.debug("WebSocket disconnected: %s", exc)
                if self._on_disconnected:
                    self._on_disconnected()

            if not self._running:
                break

            # Exponential backoff
            log.debug("Reconnecting in %.1fs...", self._reconnect_delay)
            await asyncio.sleep(self._reconnect_delay)
            self._reconnect_delay = min(
                self._reconnect_delay * 2, self._max_reconnect_delay
            )

    async def _connect_and_listen(self) -> None:
        """Single connection lifecycle."""
        url = self._ws_url()
        headers: dict[str, str] = {}

        creds = read_credentials()
        if creds and creds.token:
            headers["Authorization"] = f"Bearer {creds.token}"

        async with ws_client.connect(
            url,
            additional_headers=headers,
            ping_interval=20,
            ping_timeout=10,
        ) as ws:
            self._ws = ws
            self._reconnect_delay = 1.0  # Reset on successful connect

            # Subscribe to agent events
            await ws.send(
                json.dumps({"type": "subscribe", "agentId": self._agent_id})
            )

            if self._on_connected:
                self._on_connected()

            async for raw in ws:
                if not self._running:
                    break
                self._handle_message(raw)

        self._ws = None

    def _handle_message(self, raw: str | bytes) -> None:
        """Parse and dispatch a WebSocket message."""
        try:
            if isinstance(raw, bytes):
                raw = raw.decode("utf-8")
            msg = json.loads(raw)
        except (json.JSONDecodeError, UnicodeDecodeError):
            log.warning("Invalid WebSocket message: %r", raw[:100])
            return

        msg_type = msg.get("type", "")
        try:
            event_type = WSEventType(msg_type)
        except ValueError:
            log.debug("Unknown WebSocket message type: %s", msg_type)
            return

        # Pass the full message as event data. Individual handlers
        # extract what they need. Some messages have a "data" field
        # (location_update, mystery_clue) and others don't (agent_state_update).
        event = WSEvent(event_type, msg)

        if self._on_event:
            self._on_event(event)

    async def send_ping(self) -> None:
        """Send a manual ping to keep the connection alive."""
        if self._ws:
            await self._ws.send(json.dumps({"type": "ping"}))

    async def close(self) -> None:
        """Gracefully close the WebSocket connection."""
        self._running = False
        if self._ws:
            await self._ws.close()
            self._ws = None
