kis_bot/kis_trader/ws/program_cache.py

"""
kis_trader/ws/program_cache.py — 키움 0w 종목프로그램매매 RAM 캐시
================================================================
키움 OpenAPI+ FID (종목프로그램매매 0w, KOA Studio 기준):
  202 매수수량, 204 매도수량, 206 순매수수량(부호), 207 순매수대금(부호)
  203 매수대금, 205 매도대금 (보조)
"""
from __future__ import annotations

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


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


def _signed_int(v: Any) -> int:
    try:
        s = str(v or "0").replace(",", "").strip()
        if not s or s in ("-", "+"):
            return 0
        return int(float(s))
    except (TypeError, ValueError):
        return 0


@dataclass
class ProgramSnapshot:
    """종목 1개 당일 프로그램매매 스냅샷."""
    code: str
    buy_qty: int = 0
    sell_qty: int = 0
    net_qty: int = 0
    buy_amt: int = 0
    sell_amt: int = 0
    net_amt: int = 0
    prev_net_qty: int = 0
    delta_net_qty: int = 0
    ts: float = 0.0
    source: str = "kiwoom_0w"
    snap_time: str = ""  # YYYYMMDDHHMMSS — DB·판정 스냅 재생용

    def is_net_buy(self) -> bool:
        return self.net_qty > 0

    def is_net_sell(self) -> bool:
        return self.net_qty < 0

    def sell_buy_qty_ratio(self) -> float:
        if self.buy_qty <= 0:
            return 999.0 if self.sell_qty > 0 else 0.0
        return self.sell_qty / float(self.buy_qty)

    def to_storage_dict(self) -> Dict[str, Any]:
        """``ws_program`` INSERT·백테 복원용."""
        return {
            "code": self.code,
            "buy_qty": self.buy_qty,
            "sell_qty": self.sell_qty,
            "net_qty": self.net_qty,
            "buy_amt": self.buy_amt,
            "sell_amt": self.sell_amt,
            "net_amt": self.net_amt,
            "source": self.source,
            "ts": self.ts,
            "snap_time": self.snap_time or "",
        }


def program_snapshot_from_storage(row: Dict[str, Any]) -> ProgramSnapshot:
    """DB 행 또는 storage dict → ``ProgramSnapshot``."""
    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 ProgramSnapshot(
        code=str(row.get("code") or "").strip(),
        buy_qty=_abs_int(row.get("buy_qty")),
        sell_qty=_abs_int(row.get("sell_qty")),
        net_qty=_signed_int(row.get("net_qty")),
        buy_amt=_abs_int(row.get("buy_amt")),
        sell_amt=_abs_int(row.get("sell_amt")),
        net_amt=_signed_int(row.get("net_amt")),
        ts=ts,
        source=str(row.get("source") or "kiwoom_0w"),
        snap_time=st[:14] if st else "",
    )


def parse_kiwoom_0w_values(code: str, values: Dict[str, Any]) -> ProgramSnapshot:
    """키움 WS 0w ``values`` dict → ``ProgramSnapshot``."""
    buy_qty = _abs_int(values.get("202"))
    sell_qty = _abs_int(values.get("204"))
    net_qty = _signed_int(values.get("206"))
    if net_qty == 0 and (buy_qty > 0 or sell_qty > 0):
        net_qty = buy_qty - sell_qty

    buy_amt = _abs_int(values.get("203"))
    sell_amt = _abs_int(values.get("205"))
    net_amt = _signed_int(values.get("207"))
    if net_amt == 0 and (buy_amt > 0 or sell_amt > 0):
        net_amt = buy_amt - sell_amt

    return ProgramSnapshot(
        code=str(code).strip(),
        buy_qty=buy_qty,
        sell_qty=sell_qty,
        net_qty=net_qty,
        buy_amt=buy_amt,
        sell_amt=sell_amt,
        net_amt=net_amt,
        ts=time.time(),
        source="kiwoom_0w",
    )


class ProgramCache:
    """스레드 세이프 종목별 프로그램매매 캐시."""

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

    def update_from_kiwoom_0w(self, code: str, values: Dict[str, Any]) -> ProgramSnapshot:
        snap = parse_kiwoom_0w_values(code, values)
        with self._lock:
            prev = self._data.get(snap.code)
            if prev is not None:
                snap.prev_net_qty = prev.net_qty
                snap.delta_net_qty = snap.net_qty - prev.net_qty
            self._data[snap.code] = snap
        return snap

    def get(self, code: str, max_age_sec: float = 30.0) -> Optional[ProgramSnapshot]:
        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 remove(self, code: str) -> None:
        c = str(code or "").strip()
        with self._lock:
            self._data.pop(c, None)