"""StreamModule — real-time WebSocket feed for NAV updates and events."""

from __future__ import annotations

import asyncio
import json
import logging
from collections import defaultdict
from typing import Any, Callable, Coroutine, Dict, List, Optional

import websockets
import websockets.client

logger = logging.getLogger("openhedge.stream")


class StreamModule:
    """WebSocket client for the OpenHedge real-time feed.

    Usage::

        stream = agent.stream

        @stream.on("vault_nav_update")
        async def handle_nav(data):
            print(data)

        await stream.connect()
        await stream.subscribe(vaults=["0x..."])
        await stream.listen()
    """

    def __init__(
        self,
        *,
        feed_url: str,
        api_key: str,
        reconnect_delay: float = 2.0,
        max_reconnect_delay: float = 60.0,
    ) -> None:
        self._feed_url = feed_url
        self._api_key = api_key
        self._reconnect_delay = reconnect_delay
        self._max_reconnect_delay = max_reconnect_delay
        self._ws: Optional[websockets.client.WebSocketClientProtocol] = None
        self._handlers: Dict[str, List[Callable[..., Coroutine[Any, Any, None]]]] = defaultdict(
            list
        )
        self._connected = False
        self._current_delay = reconnect_delay

    # ------------------------------------------------------------------
    # Connection lifecycle
    # ------------------------------------------------------------------

    async def connect(self) -> None:
        """Establish the WebSocket connection.

        Waits for the server's ``connected`` acknowledgement.
        """
        self._ws = await websockets.connect(  # type: ignore[attr-defined]
            self._feed_url,
        )
        # Send auth as in-band message (not via headers)
        await self._ws.send(json.dumps({"type": "auth", "apiKey": self._api_key}))
        # Wait for "connected" message
        raw = await self._ws.recv()
        msg = json.loads(raw)
        if msg.get("type") == "connected":
            logger.info("WebSocket connected to %s", self._feed_url)
            self._connected = True
            self._current_delay = self._reconnect_delay
        else:
            logger.warning("Unexpected initial message: %s", msg)

    async def disconnect(self) -> None:
        """Gracefully close the WebSocket."""
        if self._ws:
            await self._ws.close()
            self._ws = None
            self._connected = False
            logger.info("WebSocket disconnected")

    # ------------------------------------------------------------------
    # Event handling
    # ------------------------------------------------------------------

    def on(self, event: str) -> Callable:
        """Decorator to register an async handler for *event*.

        Example::

            @stream.on("vault_nav_update")
            async def on_nav(data):
                ...
        """

        def decorator(func: Callable[..., Coroutine[Any, Any, None]]) -> Callable:
            self._handlers[event].append(func)
            return func

        return decorator

    # ------------------------------------------------------------------
    # Subscriptions
    # ------------------------------------------------------------------

    async def subscribe(self, vaults: Optional[List[str]] = None) -> None:
        """Subscribe to updates, optionally filtered by vault addresses."""
        if self._ws is None:
            raise RuntimeError("Not connected — call connect() first")
        payload = {"type": "subscribe", "vaults": vaults or []}
        await self._ws.send(json.dumps(payload))
        logger.info("Subscribed to vaults: %s", vaults or "all")

    # ------------------------------------------------------------------
    # Main event loop
    # ------------------------------------------------------------------

    async def listen(self) -> None:
        """Block and dispatch incoming events.

        Automatically reconnects on connection loss with exponential backoff.
        """
        while True:
            try:
                if self._ws is None:
                    await self.connect()

                async for raw_message in self._ws:  # type: ignore[union-attr]
                    try:
                        data = json.loads(raw_message)
                    except json.JSONDecodeError:
                        logger.warning("Non-JSON message: %s", raw_message[:200])
                        continue

                    event_type = data.get("type", "unknown")
                    event_data = data.get("data")
                    handlers = self._handlers.get(event_type, [])

                    for handler in handlers:
                        try:
                            await handler(event_data)
                        except Exception:
                            logger.exception("Handler error for event '%s'", event_type)

            except websockets.ConnectionClosed as exc:
                logger.warning("Connection closed (%s). Reconnecting...", exc)
                self._connected = False
                self._ws = None
                await self._backoff()

            except Exception as exc:
                logger.exception("Stream error: %s", exc)
                self._connected = False
                self._ws = None
                await self._backoff()

    # ------------------------------------------------------------------
    # Internals
    # ------------------------------------------------------------------

    async def _backoff(self) -> None:
        """Sleep with exponential backoff before reconnecting."""
        logger.info("Waiting %.1fs before reconnect...", self._current_delay)
        await asyncio.sleep(self._current_delay)
        self._current_delay = min(self._current_delay * 2, self._max_reconnect_delay)
