polymarket-insider-tracker/src/polymarket_insider_tracker/ingestor/websocket.py

"""WebSocket client for streaming Polymarket trade events."""

import asyncio
import json
import logging
import time
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from enum import Enum
from typing import Any

from websockets.asyncio.client import ClientConnection
from websockets.asyncio.client import connect as ws_connect
from websockets.exceptions import ConnectionClosed

from polymarket_insider_tracker.ingestor.models import TradeEvent

logger = logging.getLogger(__name__)

# Constants
DEFAULT_WS_HOST = "wss://ws-live-data.polymarket.com"
DEFAULT_PING_INTERVAL = 30  # seconds
DEFAULT_MAX_RECONNECT_DELAY = 30  # seconds
DEFAULT_INITIAL_RECONNECT_DELAY = 1  # seconds


class ConnectionState(Enum):
    """WebSocket connection states."""

    DISCONNECTED = "disconnected"
    CONNECTING = "connecting"
    CONNECTED = "connected"
    RECONNECTING = "reconnecting"


@dataclass
class StreamStats:
    """Statistics about the trade stream."""

    trades_received: int = 0
    reconnect_count: int = 0
    last_trade_time: float | None = None
    connected_since: float | None = None
    last_error: str | None = None


class TradeStreamError(Exception):
    """Base exception for trade stream errors."""


class ConnectionError(TradeStreamError):
    """Raised when connection to WebSocket fails."""


TradeCallback = Callable[[TradeEvent], Awaitable[None]]
StateCallback = Callable[[ConnectionState], Awaitable[None]]


class TradeStreamHandler:
    """WebSocket client for streaming Polymarket trade events.

    This handler maintains a persistent connection to Polymarket's real-time
    trade feed, automatically reconnecting on disconnection with exponential
    backoff.

    Example:
        >>> async def on_trade(trade: TradeEvent):
        ...     print(f"Trade: {trade.side} {trade.size} @ {trade.price}")
        ...
        >>> handler = TradeStreamHandler(on_trade=on_trade)
        >>> await handler.start()  # Blocks until stop() is called

    Attributes:
        state: Current connection state.
        stats: Statistics about the stream (trades received, reconnects, etc.).
    """

    def __init__(
        self,
        on_trade: TradeCallback,
        *,
        host: str = DEFAULT_WS_HOST,
        on_state_change: StateCallback | None = None,
        ping_interval: int = DEFAULT_PING_INTERVAL,
        max_reconnect_delay: int = DEFAULT_MAX_RECONNECT_DELAY,
        initial_reconnect_delay: int = DEFAULT_INITIAL_RECONNECT_DELAY,
        event_filter: str | None = None,
        market_filter: str | None = None,
    ) -> None:
        """Initialize the trade stream handler.

        Args:
            on_trade: Async callback invoked for each trade event.
            host: WebSocket endpoint URL.
            on_state_change: Optional callback for connection state changes.
            ping_interval: Seconds between heartbeat pings.
            max_reconnect_delay: Maximum delay between reconnection attempts.
            initial_reconnect_delay: Initial delay for reconnection backoff.
            event_filter: Optional event slug to filter trades by event.
            market_filter: Optional market slug to filter trades by market.
        """
        self._on_trade = on_trade
        self._on_state_change = on_state_change
        self._host = host
        self._ping_interval = ping_interval
        self._max_reconnect_delay = max_reconnect_delay
        self._initial_reconnect_delay = initial_reconnect_delay
        self._event_filter = event_filter
        self._market_filter = market_filter

        self._state = ConnectionState.DISCONNECTED
        self._stats = StreamStats()
        self._ws: ClientConnection | None = None
        self._running = False
        self._stop_event: asyncio.Event | None = None

    @property
    def state(self) -> ConnectionState:
        """Current connection state."""
        return self._state

    @property
    def stats(self) -> StreamStats:
        """Stream statistics."""
        return self._stats

    async def _set_state(self, new_state: ConnectionState) -> None:
        """Update state and notify callback."""
        if self._state != new_state:
            old_state = self._state
            self._state = new_state
            logger.info("Connection state: %s -> %s", old_state.value, new_state.value)

            if self._on_state_change:
                try:
                    await self._on_state_change(new_state)
                except Exception as e:
                    logger.error("Error in state change callback: %s", e)

    def _build_subscription_message(self) -> dict[str, Any]:
        """Build the WebSocket subscription message."""
        subscription: dict[str, Any] = {
            "topic": "activity",
            "type": "trades",
        }

        # Add filters if specified
        if self._event_filter:
            subscription["filters"] = json.dumps({"event_slug": self._event_filter})
        elif self._market_filter:
            subscription["filters"] = json.dumps({"market_slug": self._market_filter})

        return {"action": "subscribe", "subscriptions": [subscription]}

    async def _connect(self) -> ClientConnection:
        """Establish WebSocket connection."""
        await self._set_state(ConnectionState.CONNECTING)

        try:
            ws = await ws_connect(
                self._host,
                ping_interval=self._ping_interval,
                ping_timeout=self._ping_interval * 2,
            )

            # Send subscription message
            subscribe_msg = self._build_subscription_message()
            await ws.send(json.dumps(subscribe_msg))

            logger.info("Connected to %s and subscribed to trades", self._host)
            await self._set_state(ConnectionState.CONNECTED)
            self._stats.connected_since = time.time()

            return ws

        except Exception as e:
            logger.error("Failed to connect: %s", e)
            self._stats.last_error = str(e)
            raise ConnectionError(f"Failed to connect to {self._host}: {e}") from e

    async def _handle_message(self, message: str) -> None:
        """Parse and process an incoming WebSocket message."""
        try:
            data = json.loads(message)

            # ws-live-data pushes {connection_id, payload:{...trade fields}}
            payload = data.get("payload")
            if (
                isinstance(payload, dict)
                and "transactionHash" in payload
                and "proxyWallet" in payload
            ):
                trade = TradeEvent.from_websocket_message(payload)

                self._stats.trades_received += 1
                self._stats.last_trade_time = time.time()

                logger.debug(
                    "Trade: %s %s @ %s on %s",
                    trade.side,
                    trade.size,
                    trade.price,
                    trade.market_slug,
                )

                try:
                    await self._on_trade(trade)
                except Exception as e:
                    logger.error("Error in trade callback: %s", e)

            else:
                logger.debug("Received non-trade message: %s", str(data)[:120])

        except json.JSONDecodeError as e:
            logger.warning("Invalid JSON message: %s", e)
        except Exception as e:
            logger.error("Error processing message: %s", e)

    async def _listen(self, ws: ClientConnection) -> None:
        """Listen for messages on the WebSocket."""
        try:
            async for message in ws:
                if not self._running:
                    break

                if isinstance(message, str):
                    await self._handle_message(message)
                else:
                    logger.debug("Received binary message (%d bytes)", len(message))

        except ConnectionClosed as e:
            logger.warning("Connection closed: %s", e)
            raise
        except Exception as e:
            logger.error("Error in message loop: %s", e)
            raise

    async def _reconnect_loop(self) -> None:
        """Handle reconnection with exponential backoff."""
        delay = self._initial_reconnect_delay

        while self._running:
            try:
                await self._set_state(ConnectionState.RECONNECTING)

                logger.info("Reconnecting in %.1f seconds...", delay)
                await asyncio.sleep(delay)

                if not self._running:
                    break

                self._ws = await self._connect()
                self._stats.reconnect_count += 1
                delay = self._initial_reconnect_delay  # Reset delay on success
                return

            except Exception as e:
                logger.error("Reconnection failed: %s", e)
                self._stats.last_error = str(e)

                # Exponential backoff with jitter
                delay = min(delay * 2, self._max_reconnect_delay)

    async def start(self) -> None:
        """Connect and begin streaming trades.

        This method blocks until stop() is called. It automatically
        handles reconnection on disconnection.

        Raises:
            ConnectionError: If initial connection fails.
        """
        if self._running:
            logger.warning("Handler already running")
            return

        self._running = True
        self._stop_event = asyncio.Event()

        try:
            # Initial connection
            self._ws = await self._connect()

            # Main loop
            while self._running:
                try:
                    await self._listen(self._ws)
                except (ConnectionClosed, Exception) as e:
                    if not self._running:
                        break

                    logger.warning("Connection lost: %s", e)
                    await self._set_state(ConnectionState.DISCONNECTED)

                    # Attempt reconnection
                    await self._reconnect_loop()

                    if not self._running or self._ws is None:
                        break

        finally:
            await self._cleanup()

    async def stop(self) -> None:
        """Gracefully disconnect from the WebSocket.

        This signals the handler to stop and cleanly close the connection.
        """
        if not self._running:
            return

        logger.info("Stopping trade stream handler...")
        self._running = False

        if self._stop_event:
            self._stop_event.set()

        await self._cleanup()

    async def _cleanup(self) -> None:
        """Clean up resources."""
        if self._ws:
            try:
                await self._ws.close()
            except Exception as e:
                logger.debug("Error closing WebSocket: %s", e)
            finally:
                self._ws = None

        await self._set_state(ConnectionState.DISCONNECTED)
        logger.info("Trade stream handler stopped")

    async def __aenter__(self) -> "TradeStreamHandler":
        """Async context manager entry."""
        return self

    async def __aexit__(self, *args: Any) -> None:
        """Async context manager exit."""
        await self.stop()