kis_bot/kis_trader/ws/tick_recorder.py

"""
kis_trader/ws/tick_recorder.py — 실시간 체결 틱 RAM 링버퍼 + ws_ticks 배치 저장 (C안)
==================================================================================
[트랙 1] 종목별 deque 링버퍼 — B안 봉 내 돌파·get_recent_ticks() 용 (DB 대기 없음)
[트랙 2] 제한 큐 → 백그라운드 스레드 배치 INSERT (journalctl 파싱 아님)

해외 확장: market / exchange / currency / session 컬럼 (국내 기본 KR/KRX/KRW).
- 국내 → ``ws_ticks`` (메인 TradeDB 커넥션)
- 해외 US → ``ws_ticks_us`` (전용 raw 커넥션 — 앱 TradeDB 락과 분리)
"""
from __future__ import annotations

import datetime as dt
import queue
import threading
import time
from collections import deque
from typing import Any, Deque, Dict, List, Optional, Set

from ..utils.env import get_env_bool, get_env_float, get_env_from_db, get_env_int
from ..utils.logger import get_logger

logger = get_logger("kis_trader.tick_recorder")


class TickRecorder:
    """
    WebSocket 체결 틱 기록기.

    - ``set_record_codes()`` 로 후보(또는 구독) 종목만 필터.
    - RAM: 종목당 최대 N틱 (링버퍼).
    - DB: BATCH_SIZE / FLUSH_SEC 마다 append INSERT
      (KR→ws_ticks, US→ws_ticks_us).
    """

    def __init__(self, db=None) -> None:
        self.db = db
        self._enabled = get_env_bool("WS_TICK_SAVE_ENABLED", True)
        self._buffer_max = get_env_int("WS_TICK_BUFFER_MAX_PER_CODE", 500)
        self._batch_size = get_env_int("WS_TICK_DB_BATCH_SIZE", 200)
        self._flush_sec = get_env_float("WS_TICK_DB_FLUSH_SEC", 3.0)
        self._queue_max = get_env_int("WS_TICK_WRITE_QUEUE_MAX", 50000)
        self._keep_days = get_env_int("WS_TICK_KEEP_DAYS", 7)
        self._default_market = (get_env_from_db("WS_TICK_DEFAULT_MARKET", "KR") or "KR").strip().upper()
        self._default_exchange = (get_env_from_db("WS_TICK_DEFAULT_EXCHANGE", "KRX") or "KRX").strip()
        self._default_currency = (get_env_from_db("WS_TICK_DEFAULT_CURRENCY", "KRW") or "KRW").strip().upper()

        self._lock = threading.Lock()
        self._buffers: Dict[str, Deque[dict]] = {}
        self._record_codes: Optional[Set[str]] = None
        self._write_queue: queue.Queue = queue.Queue(maxsize=max(1000, self._queue_max))
        self._writer_thread: Optional[threading.Thread] = None
        self._last_cleanup_ts = 0.0
        self._drops = 0
        # 해외 틱 전용 DB 커넥션 (메인 TradeDB._lock 과 분리)
        self._us_write_conn: Any = None

        if self._enabled and self.db is not None:
            self._ensure_table()
            self._start_writer()
            logger.info(
                "✅ TickRecorder 활성 (buf=%d/종목, batch=%d, flush=%.1fs, keep=%dd, us→ws_ticks_us)",
                self._buffer_max, self._batch_size, self._flush_sec, self._keep_days,
            )
        elif self._enabled and self.db is None:
            logger.warning("⚠️ TickRecorder: db=None → RAM 링버퍼만 (DB 저장 비활성)")
        else:
            logger.info("ℹ️ TickRecorder 비활성 (WS_TICK_SAVE_ENABLED=false)")

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def set_record_codes(self, codes: Optional[Set[str]]) -> None:
        """None=필터 없음(구독 전체), 빈 set=저장 안 함."""
        with self._lock:
            if codes is None:
                self._record_codes = None
            else:
                self._record_codes = {str(c).strip() for c in codes if c}

    def on_tick(
        self,
        code: str,
        price: float,
        volume: int,
        tick_time: str,
        *,
        market: Optional[str] = None,
        exchange: Optional[str] = None,
        currency: Optional[str] = None,
        source: str = "kis",
        session: Optional[str] = None,
        tick_seq: Optional[int] = None,
        persist_db: bool = True,
        cntr_str: Optional[float] = None,
        upper_limit_time: Optional[str] = None,
        tick_time_raw: Optional[str] = None,
    ) -> None:
        """체결 틱 1건 — WS 핫패스에서 호출 (논블로킹).

        persist_db=False 이면 RAM 링버퍼만 갱신 (DB INSERT 생략).
        LS 는 ``ls_ws_ticks`` 에 이미 적재하므로 ``ws_ticks`` 이중 저장 방지용.
        """
        if not self._enabled or price <= 0:
            return
        code = (code or "").strip()
        if not code:
            return
        mkt = (market or self._default_market).upper()[:8]
        # 국내: WS_TICK_RECORD_SCOPE(candidates 등) 필터 적용.
        # 해외(US): 영구구독이 국내 후보 집합에 없어 필터에 걸리면
        # 봉(ws_candles)은 쌓이는데 틱(ws_ticks_us)만 끊기는 사고가 남.
        # → US 는 필터 우회 (저장 ON/OFF 는 해외 WS 의 WS_TICK_SAVE_US_ENABLED).
        if mkt != "US":
            with self._lock:
                filt = self._record_codes
            if filt is not None and code not in filt:
                return

        tick_time_full = self._normalize_tick_time(tick_time)
        recv_ts = dt.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        item = {
            "market": mkt,
            "exchange": (exchange or self._default_exchange)[:16] if (exchange or self._default_exchange) else None,
            "code": code[:32],
            "tick_time": tick_time_full,
            "price": float(price),
            "volume": int(volume or 0),
            "tick_seq": tick_seq,
            "session": (session or "")[:8] or None,
            "currency": (currency or self._default_currency).upper()[:8],
            "source": (source or "kis")[:16],
            "recv_ts": recv_ts,
        }
        if cntr_str is not None:
            item["cntr_str"] = float(cntr_str)
        if upper_limit_time:
            item["upper_limit_time"] = str(upper_limit_time).strip()[:8]
        if tick_time_raw:
            # WS 원본 체결시각 — 자르지 않음 (초·밀리초·영업일 포함 그대로)
            item["tick_time_raw"] = str(tick_time_raw).strip()[:64]

        with self._lock:
            buf = self._buffers.get(code)
            if buf is None:
                buf = deque(maxlen=max(10, self._buffer_max))
                self._buffers[code] = buf
            buf.append(item)

        # LS 등: 전용 테이블에 이미 쓰면 ws_ticks 중복 INSERT 생략
        if not persist_db or self.db is None:
            return
        try:
            self._write_queue.put_nowait(item)
        except queue.Full:
            self._drops += 1
            if self._drops == 1 or self._drops % 5000 == 0:
                logger.warning(
                    "⚠️ TickRecorder 쓰기 Queue 가득참 — 틱 DROP 누적=%d (code=%s)",
                    self._drops, code,
                )

    def get_recent_ticks(self, code: str, limit: int = 100) -> List[dict]:
        """RAM 링버퍼 최근 틱 (오래된→최신 순)."""
        code = (code or "").strip()
        if not code:
            return []
        with self._lock:
            buf = self._buffers.get(code)
            if not buf:
                return []
            rows = list(buf)
        n = max(1, int(limit))
        return rows[-n:]

    def get_last_price(self, code: str) -> Optional[float]:
        ticks = self.get_recent_ticks(code, limit=1)
        if not ticks:
            return None
        p = ticks[-1].get("price")
        return float(p) if p and float(p) > 0 else None

    def remove_code(self, code: str) -> None:
        """구독 해제 시 RAM 버퍼 정리."""
        code = (code or "").strip()
        if not code:
            return
        with self._lock:
            self._buffers.pop(code, None)

    # ------------------------------------------------------------------
    # 내부
    # ------------------------------------------------------------------

    @staticmethod
    def _normalize_tick_time(tick_time: str) -> str:
        """HHMMSS · YYYYMMDDHHMMSS · 그 이상(밀리초 등) → tick_time 컬럼용 14자리+."""
        raw = (tick_time or "").strip().replace(":", "").replace("-", "").replace(" ", "")
        if len(raw) >= 14 and raw[:14].isdigit():
            return raw[:14]
        if len(raw) == 6 and raw.isdigit():
            today = dt.datetime.now().strftime("%Y%m%d")
            return today + raw
        if len(raw) > 6 and raw.isdigit():
            # KIS BSOP_DATE+STCK_CNTG_HOUR 등 — 앞 14자리가 유효하면 사용
            if len(raw) >= 14:
                return raw[:14]
        return dt.datetime.now().strftime("%Y%m%d%H%M%S")

    def _ensure_table(self) -> None:
        try:
            if hasattr(self.db, "ensure_ws_ticks_table"):
                self.db.ensure_ws_ticks_table()
            if hasattr(self.db, "ensure_ws_ticks_us_table"):
                self.db.ensure_ws_ticks_us_table()
        except Exception as e:
            logger.warning("ws_ticks/_us 테이블 확인 실패: %s", e)
        # 해외 INSERT 전용 커넥션 (메인 전략 SELECT 와 파이썬 락 비공유)
        try:
            if hasattr(self.db, "new_raw_connection"):
                self._us_write_conn = self.db.new_raw_connection()
                logger.info("✅ TickRecorder US writer 전용 커넥션 (ws_ticks_us)")
        except Exception as e:
            self._us_write_conn = None
            logger.warning("US 전용 커넥션 실패 — 메인 DB로 폴백: %s", e)

    def _start_writer(self) -> None:
        self._writer_thread = threading.Thread(
            target=self._writer_loop,
            name="TickDBWriter",
            daemon=True,
        )
        self._writer_thread.start()

    def _writer_loop(self) -> None:
        batch: List[dict] = []
        last_flush = time.time()
        while True:
            try:
                timeout = max(0.1, self._flush_sec - (time.time() - last_flush))
                item = self._write_queue.get(timeout=timeout)
                if item is None:
                    break
                batch.append(item)
                self._write_queue.task_done()
            except queue.Empty:
                pass

            now = time.time()
            should = len(batch) >= self._batch_size or (
                batch and (now - last_flush) >= self._flush_sec
            )
            if should and batch:
                self._flush_batch(batch)
                batch = []
                last_flush = now
            if now - self._last_cleanup_ts >= 3600.0:
                self._maybe_cleanup_old()
                self._last_cleanup_ts = now

        if batch:
            self._flush_batch(batch)

    def _flush_batch(self, batch: List[dict]) -> None:
        if not batch or not self.db:
            return
        us_batch = [x for x in batch if str(x.get("market") or "").upper() == "US"]
        kr_batch = [x for x in batch if str(x.get("market") or "").upper() != "US"]
        try:
            if kr_batch and hasattr(self.db, "insert_ws_ticks_batch"):
                self.db.insert_ws_ticks_batch(kr_batch)
                logger.debug("💾 [틱배치] %d건 → ws_ticks", len(kr_batch))
        except Exception as e:
            logger.warning("ws_ticks 배치 INSERT 실패 (%d건): %s", len(kr_batch), e)
        try:
            if us_batch and hasattr(self.db, "insert_ws_ticks_us_batch"):
                self.db.insert_ws_ticks_us_batch(
                    us_batch, conn=self._us_write_conn,
                )
                logger.debug("💾 [틱배치] %d건 → ws_ticks_us", len(us_batch))
            elif us_batch and hasattr(self.db, "insert_ws_ticks_batch"):
                # 구버전 DB 헬퍼 폴백 (분리 전)
                self.db.insert_ws_ticks_batch(us_batch)
                logger.debug("💾 [틱배치] %d건 → ws_ticks(폴백 US)", len(us_batch))
        except Exception as e:
            logger.warning("ws_ticks_us 배치 INSERT 실패 (%d건): %s", len(us_batch), e)

    def _maybe_cleanup_old(self) -> None:
        if self._keep_days <= 0:
            return
        try:
            if hasattr(self.db, "cleanup_old_ws_ticks"):
                self.db.cleanup_old_ws_ticks(keep_days=self._keep_days)
        except Exception as e:
            logger.debug("ws_ticks 정리 실패: %s", e)
        try:
            if hasattr(self.db, "cleanup_old_ws_ticks_us"):
                self.db.cleanup_old_ws_ticks_us(keep_days=self._keep_days)
        except Exception as e:
            logger.debug("ws_ticks_us 정리 실패: %s", e)