kis_bot/kis_trader/ws/orderbook_cache.py

"""
kis_trader/ws/orderbook_cache.py — 키움 0D 호가잔량 RAM 캐시
================================================================
키움 OpenAPI+ FID (주식호가잔량 0D):
  매도호가 1~10: 41~50, 매수호가 1~10: 51~60
  매도수량 1~10: 61~70, 매수수량 1~10: 71~80
  매도총잔량: 121, 매수총잔량: 125
"""
from __future__ import annotations

import json
import threading
import time
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple


def _abs_int(v: Any) -> int:
    try:
        return int(abs(float(str(v or "0").replace(",", ""))))
    except (TypeError, ValueError):
        return 0


@dataclass
class OrderbookLevel:
    price: int = 0
    qty: int = 0


@dataclass
class OrderbookSnapshot:
    """종목 1개 호가 스냅샷 (매수호가=비싼 순, 매도호가=싼 순)."""
    code: str
    bids: List[OrderbookLevel] = field(default_factory=list)
    asks: List[OrderbookLevel] = field(default_factory=list)
    total_bid_qty: int = 0
    total_ask_qty: int = 0
    ts: float = 0.0
    source: str = "kiwoom_0d"
    snap_time: str = ""  # YYYYMMDDHHMMSS — DB·판정 스냅 재생용

    def best_bid(self) -> int:
        return self.bids[0].price if self.bids else 0

    def best_ask(self) -> int:
        return self.asks[0].price if self.asks else 0

    def spread_pct(self) -> float:
        bb, ba = self.best_bid(), self.best_ask()
        if bb <= 0 or ba <= 0:
            return 999.0
        mid = (bb + ba) / 2.0
        if mid <= 0:
            return 999.0
        return (ba - bb) / mid * 100.0

    def bid_qty_sum(self, levels: int = 3) -> int:
        n = max(1, int(levels))
        return sum(lv.qty for lv in self.bids[:n])

    def ask_qty_sum(self, levels: int = 3) -> int:
        n = max(1, int(levels))
        return sum(lv.qty for lv in self.asks[:n])

    def to_kis_bid_dict(self, levels: int = 10) -> Dict[str, Any]:
        """``orderbook_sell.parse_kis_bid_levels`` 호환 dict."""
        out: Dict[str, Any] = {}
        for i, lv in enumerate(self.bids[: max(1, levels)], start=1):
            out[f"bidp{i}"] = lv.price
            out[f"bidp_rsqn{i}"] = lv.qty
        return out

    def to_storage_dict(self) -> Dict[str, Any]:
        """``ws_orderbook`` INSERT·백테 복원용."""
        return {
            "code": self.code,
            "best_bid": self.best_bid(),
            "best_ask": self.best_ask(),
            "total_bid_qty": self.total_bid_qty,
            "total_ask_qty": self.total_ask_qty,
            "bid_qty_l3": self.bid_qty_sum(3),
            "ask_qty_l3": self.ask_qty_sum(3),
            "levels_json": json.dumps(
                {
                    "bids": [[lv.price, lv.qty] for lv in self.bids[:10]],
                    "asks": [[lv.price, lv.qty] for lv in self.asks[:10]],
                },
                ensure_ascii=False,
                separators=(",", ":"),
            ),
            "source": self.source,
            "ts": self.ts,
            "snap_time": self.snap_time or "",
        }


def orderbook_snapshot_from_storage(row: Dict[str, Any]) -> OrderbookSnapshot:
    """DB 행 또는 storage dict → ``OrderbookSnapshot``."""
    bids: List[OrderbookLevel] = []
    asks: List[OrderbookLevel] = []
    raw = row.get("levels_json")
    if raw:
        try:
            lv = json.loads(raw) if isinstance(raw, str) else raw
            for px, qty in (lv.get("bids") or []):
                bids.append(OrderbookLevel(price=_abs_int(px), qty=_abs_int(qty)))
            for px, qty in (lv.get("asks") or []):
                asks.append(OrderbookLevel(price=_abs_int(px), qty=_abs_int(qty)))
        except (TypeError, ValueError, json.JSONDecodeError):
            pass
    total_bid = _abs_int(row.get("total_bid_qty"))
    total_ask = _abs_int(row.get("total_ask_qty"))
    if total_bid <= 0 and bids:
        total_bid = sum(b.qty for b in bids)
    if total_ask <= 0 and asks:
        total_ask = sum(a.qty for a in asks)
    ts_raw = row.get("ts")
    try:
        ts = float(ts_raw) if ts_raw not in (None, "") else time.time()
    except (TypeError, ValueError):
        ts = time.time()
    st = str(row.get("snap_time") or "").strip()
    return OrderbookSnapshot(
        code=str(row.get("code") or "").strip(),
        bids=bids,
        asks=asks,
        total_bid_qty=total_bid,
        total_ask_qty=total_ask,
        ts=ts,
        source=str(row.get("source") or "kiwoom_0d"),
        snap_time=st[:14] if st else "",
    )


def parse_kiwoom_0d_values(code: str, values: Dict[str, Any]) -> OrderbookSnapshot:
    """키움 WS 0D ``values`` dict → ``OrderbookSnapshot``."""
    asks: List[OrderbookLevel] = []
    bids: List[OrderbookLevel] = []
    for i in range(1, 11):
        ask_px = _abs_int(values.get(str(40 + i)))
        bid_px = _abs_int(values.get(str(50 + i)))
        ask_qty = _abs_int(values.get(str(60 + i)))
        bid_qty = _abs_int(values.get(str(70 + i)))
        if ask_px > 0:
            asks.append(OrderbookLevel(price=ask_px, qty=ask_qty))
        if bid_px > 0:
            bids.append(OrderbookLevel(price=bid_px, qty=bid_qty))
    total_ask = _abs_int(values.get("121"))
    total_bid = _abs_int(values.get("125"))
    if total_bid <= 0 and bids:
        total_bid = sum(b.qty for b in bids)
    if total_ask <= 0 and asks:
        total_ask = sum(a.qty for a in asks)
    return OrderbookSnapshot(
        code=str(code).strip(),
        asks=asks,
        bids=bids,
        total_bid_qty=total_bid,
        total_ask_qty=total_ask,
        ts=time.time(),
        source="kiwoom_0d",
    )


class OrderbookCache:
    """스레드 세이프 종목별 호가 캐시."""

    def __init__(self) -> None:
        self._data: Dict[str, OrderbookSnapshot] = {}
        self._lock = threading.Lock()

    def update_from_kiwoom_0d(self, code: str, values: Dict[str, Any]) -> OrderbookSnapshot:
        snap = parse_kiwoom_0d_values(code, values)
        with self._lock:
            self._data[snap.code] = snap
        return snap

    def get(self, code: str, max_age_sec: float = 3.0) -> Optional[OrderbookSnapshot]:
        c = str(code or "").strip()
        if not c:
            return None
        with self._lock:
            snap = self._data.get(c)
        if not snap:
            return None
        if max_age_sec > 0 and (time.time() - snap.ts) > max_age_sec:
            return None
        return snap

    def get_kis_bid_dict(self, code: str, max_age_sec: float = 3.0) -> Optional[Dict[str, Any]]:
        snap = self.get(code, max_age_sec=max_age_sec)
        if not snap or not snap.bids:
            return None
        return snap.to_kis_bid_dict()

    def remove(self, code: str) -> None:
        c = str(code or "").strip()
        with self._lock:
            self._data.pop(c, None)