data_feed.py

"""Binance spot WebSocket feed: aggTrade ticks + kline closes (Chapter 3).

Uses python-binance ``ThreadedWebsocketManager`` against real spot streams
``wss://stream.binance.com/ws/<symbol>@aggTrade`` and ``@kline_<interval>``.
Reconnect backoff is aligned to 1s, 2s, 4s, … capped at 30s with loguru.
"""

from __future__ import annotations

import threading
import time
from collections import deque
from datetime import datetime, timezone
from typing import Callable, Deque, Optional

from loguru import logger

from binance.ws.reconnecting_websocket import ReconnectingWebsocket
from binance.ws.streams import BinanceSocketManager, ThreadedWebsocketManager

from config import Config
from microstructure import TradeEvent
from models.candle import Candle
from models.order_book import OrderBookSnapshot

OnTick = Callable[[float, datetime], None]
OnCandleClose = Callable[[Candle], None]
OnOrderBook = Callable[["OrderBookSnapshot"], None]
OnTrade = Callable[["TradeEvent"], None]

_BINANCE_SPOT_WS_BASE = "wss://stream.binance.com/"
_reconnect_patch_applied = False


def _apply_binance_reconnect_backoff() -> None:
    """Tune library reconnect: many retries; backoff 1, 2, 4, … max 30s; log each wait."""

    global _reconnect_patch_applied
    if _reconnect_patch_applied:
        return

    ReconnectingWebsocket.MAX_RECONNECTS = 999_999

    def _get_reconnect_wait_patched(self: ReconnectingWebsocket, attempts: int) -> float:
        wait = float(min(30, 2 ** max(0, attempts - 1)))
        logger.warning(
            "Binance WebSocket reconnect | stream={} | attempt={} | backoff_s={}",
            getattr(self, "_path", "?"),
            attempts,
            wait,
        )
        return wait

    ReconnectingWebsocket._get_reconnect_wait = _get_reconnect_wait_patched  # type: ignore[method-assign]
    _reconnect_patch_applied = True


class _StreamTWManager(ThreadedWebsocketManager):
    """Use explicit ``wss://stream.binance.com/`` base (no ``:9443``) like official docs."""

    async def _before_socket_listener_start(self) -> None:
        await super()._before_socket_listener_start()
        if self._bsm is not None:
            self._bsm.STREAM_URL = _BINANCE_SPOT_WS_BASE


class CandleBuffer:
    """Rolling buffer of the last ``maxlen`` *closed* candles (oldest on the left)."""

    def __init__(self, maxlen: int = 200) -> None:
        if maxlen < 1:
            raise ValueError("CandleBuffer maxlen must be >= 1")
        self._maxlen = maxlen
        self._dq: Deque[Candle] = deque(maxlen=maxlen)

    def append_closed(self, candle: Candle) -> None:
        if not candle.is_closed:
            return
        self._dq.append(candle)

    def __len__(self) -> int:
        return len(self._dq)

    def get_closes(self) -> list[float]:
        """Close prices in chronological order (oldest first)."""
        return [c.close for c in self._dq]

    def get_candle(self, n: int) -> Candle:
        """Return the *n*th most recently closed candle (``n=0`` is the latest close)."""
        if n < 0 or n >= len(self._dq):
            raise IndexError(f"candle index {n} out of range (have {len(self._dq)} candles)")
        return self._dq[-1 - n]


class DataFeed:
    """Binance spot feed: ticks + kline closes + order book depth + trade flow."""

    def __init__(
        self,
        config: Config,
        on_tick: OnTick,
        on_candle_close: OnCandleClose,
        on_order_book: Optional[OnOrderBook] = None,
        on_trade: Optional[OnTrade] = None,
    ) -> None:
        self._config = config
        self._on_tick = on_tick
        self._on_candle_close = on_candle_close
        self._on_order_book = on_order_book
        self._on_trade = on_trade
        self._twm: Optional[_StreamTWManager] = None
        self._agg_path: Optional[str] = None
        self._kline_path: Optional[str] = None
        self._depth_path: Optional[str] = None
        self._lock = threading.Lock()
        self._started = False

    def start(self) -> None:
        with self._lock:
            if self._started:
                logger.warning("DataFeed.start() called while already running; ignored")
                return

            _apply_binance_reconnect_backoff()

            sym = self._config.SYMBOL.upper()
            tf = self._config.TIMEFRAME
            logger.info(
                "DataFeed: connecting to Binance spot | symbol={} | kline={} | "
                "streams=aggTrade + kline",
                sym,
                tf,
            )

            self._twm = _StreamTWManager(
                api_key=self._config.BINANCE_API_KEY or None,
                api_secret=self._config.BINANCE_API_SECRET or None,
                tld="com",
                testnet=False,
            )
            self._twm.start()

            deadline = time.monotonic() + 30.0
            while time.monotonic() < deadline:
                if self._twm._bsm is not None:
                    break
                if not self._twm._running:
                    self._twm.join(timeout=10.0)
                    self._twm = None
                    raise RuntimeError(
                        "DataFeed: Binance client failed to start (check network / DNS / api.binance.com)"
                    )
                time.sleep(0.05)
            else:
                self._twm.stop()
                self._twm.join(timeout=10.0)
                self._twm = None
                raise RuntimeError("DataFeed: Binance socket manager failed to initialize (timeout)")

            logger.success("DataFeed: Binance threaded client ready; opening WebSocket streams")

            self._agg_path = self._twm.start_aggtrade_socket(self._handle_agg_trade, sym)
            self._kline_path = self._twm.start_kline_socket(
                self._handle_kline,
                sym,
                interval=tf,
            )
            if self._on_order_book is not None or self._on_trade is not None:
                self._depth_path = self._twm.start_depth_socket(
                    self._handle_depth,
                    sym,
                    depth=20,
                )
            self._started = True
            logger.info(
                "DataFeed: streams started | aggTrade_id={} | kline_id={} | depth_id={}",
                self._agg_path,
                self._kline_path,
                self._depth_path,
            )

    def stop(self) -> None:
        with self._lock:
            if not self._started or self._twm is None:
                self._started = False
                return

            logger.info("DataFeed: shutting down WebSocket streams")
            if self._agg_path:
                self._twm.stop_socket(self._agg_path)
            if self._kline_path:
                self._twm.stop_socket(self._kline_path)
            if self._depth_path:
                self._twm.stop_socket(self._depth_path)
            self._twm.stop()
            self._twm.join(timeout=15.0)
            self._twm = None
            self._agg_path = None
            self._kline_path = None
            self._started = False
            logger.info("DataFeed: stopped cleanly")

    def _handle_agg_trade(self, msg: dict) -> None:
        try:
            if msg.get("e") == "error":
                logger.error(
                    "aggTrade stream error | type={} | message={}",
                    msg.get("type"),
                    msg.get("m"),
                )
                return
            if msg.get("e") != "aggTrade":
                return
            price = float(msg["p"])
            qty = float(msg["q"])
            ts_ms = int(msg.get("T") if msg.get("T") is not None else msg["E"])
            ts = datetime.fromtimestamp(ts_ms / 1000.0, tz=timezone.utc)
            self._on_tick(price, ts)
            if self._on_trade is not None:
                event = TradeEvent(
                    timestamp=ts,
                    price=price,
                    qty=qty,
                    is_buyer_maker=bool(msg.get("m", False)),
                )
                self._on_trade(event)
        except Exception as exc:  # noqa: BLE001
            logger.exception("aggTrade handler failed: {}", exc)

    def _handle_depth(self, msg: dict) -> None:
        try:
            if msg.get("e") == "error":
                return
            ts_ms = int(msg.get("T", msg.get("E", 0)))
            ts = datetime.fromtimestamp(ts_ms / 1000.0, tz=timezone.utc) if ts_ms else datetime.now(timezone.utc)
            # python-binance depth messages can arrive in two shapes:
            # - {"b": [...], "a": [...]} (most common)
            # - {"bids": [...], "asks": [...]} (some internal helpers)
            raw_bids = msg.get("bids")
            raw_asks = msg.get("asks")
            if raw_bids is None and raw_asks is None:
                raw_bids = msg.get("b", [])
                raw_asks = msg.get("a", [])

            bids = [(float(p), float(q)) for p, q in (raw_bids or []) if float(q) > 0]
            asks = [(float(p), float(q)) for p, q in (raw_asks or []) if float(q) > 0]
            bids.sort(key=lambda x: x[0], reverse=True)
            asks.sort(key=lambda x: x[0])
            snap = OrderBookSnapshot(timestamp=ts, bids=bids, asks=asks)
            if self._on_order_book is not None:
                self._on_order_book(snap)
        except Exception as exc:  # noqa: BLE001
            logger.exception("depth handler failed: {}", exc)

    def _handle_kline(self, msg: dict) -> None:
        try:
            if msg.get("e") == "error":
                logger.error(
                    "kline stream error | type={} | message={}",
                    msg.get("type"),
                    msg.get("m"),
                )
                return
            if msg.get("e") != "kline":
                return
            k = msg["k"]
            is_closed = bool(k["x"])
            open_ms = int(k["t"])
            cndl = Candle(
                open_time=datetime.fromtimestamp(open_ms / 1000.0, tz=timezone.utc),
                open=float(k["o"]),
                high=float(k["h"]),
                low=float(k["l"]),
                close=float(k["c"]),
                volume=float(k["v"]),
                is_closed=is_closed,
            )
            if is_closed:
                self._on_candle_close(cndl)
        except Exception as exc:  # noqa: BLE001
            logger.exception("kline handler failed: {}", exc)


if __name__ == "__main__":
    from config import get_config

    cfg = get_config()
    stop_at = time.monotonic() + 60.0

    def on_tick(price: float, ts: datetime) -> None:
        print(f"tick price={price:.2f} time={ts.isoformat()}")

    def on_candle_close(candle: Candle) -> None:
        print(
            f"candle_close ohlc=({candle.open:.2f},{candle.high:.2f},"
            f"{candle.low:.2f},{candle.close:.2f}) vol={candle.volume:.6f} "
            f"open_time={candle.open_time.isoformat()}"
        )

    feed = DataFeed(cfg, on_tick, on_candle_close)
    try:
        feed.start()
        while time.monotonic() < stop_at:
            time.sleep(0.2)
    finally:
        feed.stop()