microdao-daarion/services/senpai-md-consumer/senpai/md_consumer/state.py

"""
State management — LatestState + RollingWindow.

All structures are asyncio-safe (no locks needed — single-threaded event loop).
"""
from __future__ import annotations

import time
from collections import deque
from dataclasses import dataclass
from datetime import datetime
from typing import Optional

from senpai.md_consumer.models import QuoteEvent, TradeEvent


@dataclass
class LatestTrade:
    symbol: str
    price: float
    size: float
    side: Optional[str]
    provider: str
    ts_recv: datetime
    ts_exchange: Optional[datetime] = None


@dataclass
class LatestQuote:
    symbol: str
    bid: float
    ask: float
    bid_size: float
    ask_size: float
    provider: str
    ts_recv: datetime
    ts_exchange: Optional[datetime] = None


@dataclass
class TradeRecord:
    """Compact trade record for rolling window."""

    price: float
    size: float
    ts: float  # monotonic seconds
    ts_exchange: Optional[datetime] = None
    ts_recv: Optional[datetime] = None


@dataclass
class QuoteRecord:
    """Compact quote record for rolling window."""

    bid: float
    ask: float
    bid_size: float
    ask_size: float
    ts: float  # monotonic seconds
    ts_exchange: Optional[datetime] = None
    ts_recv: Optional[datetime] = None


class RollingWindow:
    """
    Fixed-duration rolling window using deque.

    Efficient: O(1) append, amortised O(1) eviction.
    No pandas dependency.
    """

    def __init__(self, window_seconds: float = 60.0) -> None:
        self._window = window_seconds
        self._trades: deque[TradeRecord] = deque()
        self._quotes: deque[QuoteRecord] = deque()

    def add_trade(self, trade: TradeRecord) -> None:
        self._trades.append(trade)
        self._evict_trades()

    def add_quote(self, quote: QuoteRecord) -> None:
        self._quotes.append(quote)
        self._evict_quotes()

    def _evict_trades(self) -> None:
        cutoff = time.monotonic() - self._window
        while self._trades and self._trades[0].ts < cutoff:
            self._trades.popleft()

    def _evict_quotes(self) -> None:
        cutoff = time.monotonic() - self._window
        while self._quotes and self._quotes[0].ts < cutoff:
            self._quotes.popleft()

    @property
    def trades(self) -> deque[TradeRecord]:
        self._evict_trades()
        return self._trades

    @property
    def quotes(self) -> deque[QuoteRecord]:
        self._evict_quotes()
        return self._quotes

    def trades_since(self, seconds_ago: float) -> list[TradeRecord]:
        """Return trades within the last N seconds."""
        cutoff = time.monotonic() - seconds_ago
        return [t for t in self._trades if t.ts >= cutoff]

    def quotes_since(self, seconds_ago: float) -> list[QuoteRecord]:
        """Return quotes within the last N seconds."""
        cutoff = time.monotonic() - seconds_ago
        return [q for q in self._quotes if q.ts >= cutoff]


class LatestState:
    """
    Maintains latest trade/quote per symbol + rolling windows.
    """

    def __init__(self, window_seconds: float = 60.0) -> None:
        self._window_seconds = window_seconds
        self._latest_trade: dict[str, LatestTrade] = {}
        self._latest_quote: dict[str, LatestQuote] = {}
        self._windows: dict[str, RollingWindow] = {}
        self._event_count = 0

    def _get_window(self, symbol: str) -> RollingWindow:
        if symbol not in self._windows:
            self._windows[symbol] = RollingWindow(self._window_seconds)
        return self._windows[symbol]

    def update_trade(self, event: TradeEvent) -> None:
        """Update latest trade and rolling window."""
        sym = event.symbol.upper()

        self._latest_trade[sym] = LatestTrade(
            symbol=sym,
            price=event.price,
            size=event.size,
            side=event.side,
            provider=event.provider,
            ts_recv=event.ts_recv,
            ts_exchange=event.ts_exchange,
        )

        self._get_window(sym).add_trade(
            TradeRecord(
                price=event.price,
                size=event.size,
                ts=time.monotonic(),
                ts_exchange=event.ts_exchange,
                ts_recv=event.ts_recv,
            )
        )
        self._event_count += 1

    def update_quote(self, event: QuoteEvent) -> None:
        """Update latest quote and rolling window."""
        sym = event.symbol.upper()

        self._latest_quote[sym] = LatestQuote(
            symbol=sym,
            bid=event.bid,
            ask=event.ask,
            bid_size=event.bid_size,
            ask_size=event.ask_size,
            provider=event.provider,
            ts_recv=event.ts_recv,
            ts_exchange=event.ts_exchange,
        )

        self._get_window(sym).add_quote(
            QuoteRecord(
                bid=event.bid,
                ask=event.ask,
                bid_size=event.bid_size,
                ask_size=event.ask_size,
                ts=time.monotonic(),
                ts_exchange=event.ts_exchange,
                ts_recv=event.ts_recv,
            )
        )
        self._event_count += 1

    def get_latest_trade(self, symbol: str) -> LatestTrade | None:
        return self._latest_trade.get(symbol.upper())

    def get_latest_quote(self, symbol: str) -> LatestQuote | None:
        return self._latest_quote.get(symbol.upper())

    def get_window(self, symbol: str) -> RollingWindow | None:
        return self._windows.get(symbol.upper())

    @property
    def symbols(self) -> list[str]:
        return sorted(
            set(list(self._latest_trade.keys()) + list(self._latest_quote.keys()))
        )

    @property
    def event_count(self) -> int:
        return self._event_count

    def to_dict(self, symbol: str) -> dict:
        """Serialise latest state for API."""
        sym = symbol.upper()
        result: dict = {"symbol": sym}

        trade = self._latest_trade.get(sym)
        if trade:
            result["latest_trade"] = {
                "price": trade.price,
                "size": trade.size,
                "side": trade.side,
                "provider": trade.provider,
                "ts_recv": trade.ts_recv.isoformat() if trade.ts_recv else None,
            }

        quote = self._latest_quote.get(sym)
        if quote:
            result["latest_quote"] = {
                "bid": quote.bid,
                "ask": quote.ask,
                "bid_size": quote.bid_size,
                "ask_size": quote.ask_size,
                "provider": quote.provider,
                "ts_recv": quote.ts_recv.isoformat() if quote.ts_recv else None,
            }

        window = self._windows.get(sym)
        if window:
            result["window"] = {
                "trades_count": len(window.trades),
                "quotes_count": len(window.quotes),
            }

        return result