kis_bot/kis_trader/backtest/trigger_snapshot_loader.py

"""
kis_trader/backtest/trigger_snapshot_loader.py — ws_orderbook / ws_program 백테 로더
==================================================================================
``ws_ticks`` 와 동일하게 분봉(YYYYMMDDHHMM) 키로 인덱싱 → TRIGGER 직전 params 주입.
"""
from __future__ import annotations

import logging
from collections import defaultdict
from typing import Any, Dict, List, Optional, Set, Tuple

from kis_trader.utils.env import get_env_bool, get_env_from_db
from kis_trader.ws.orderbook_cache import OrderbookSnapshot, orderbook_snapshot_from_storage
from kis_trader.ws.program_cache import ProgramSnapshot, program_snapshot_from_storage
from kis_trader.ws.trigger_eval_recorder import FILTER_EVAL_SOURCE

logger = logging.getLogger("trigger_snapshot_loader")

LOG_BACKFILL_SOURCE = "log_backfill"
# 키움 실시간 호가 주기 스냅샷(본체 levels_json 보유) — 스프레드 등 재계산용.
# log_backfill(판정만)·filter_eval(판정시점만)과 달리 종목·시각을 폭넓게 커버한다.
KIWOOM_BODY_SOURCE = "kiwoom_0d"


def backtest_use_kiwoom_body_snapshot(
    params: Optional[Dict[str, Any]] = None,
    *,
    strategy: str = "",
) -> bool:
    """백테에서 kiwoom_0d 본체로 호가필터를 **재계산**할지 여부 (기본 false).

    기본 false: 기존처럼 log_backfill/filter_eval 판정 재생(실매 정합).
    true 일 때만 본체를 적재해 스프레드 등 임계값을 바꿔 다시 계산한다.
    웹 스프레드 입력칸·파람서치 스프레드 축이 켜질 때 params 로 true 주입.
    """
    if params is not None and params.get("backtest_use_kiwoom_body_snapshot") is not None:
        return bool(params.get("backtest_use_kiwoom_body_snapshot"))
    if params is not None and params.get("_backtest_use_kiwoom_body") is not None:
        return bool(params.get("_backtest_use_kiwoom_body"))
    pfx = (strategy or "").strip().upper()
    if pfx == "MOMENTUM":
        sk = "MOMENTUM_BACKTEST_USE_KIWOOM_BODY_SNAPSHOT"
    elif pfx == "BREAKOUT":
        sk = "BREAKOUT_BACKTEST_USE_KIWOOM_BODY_SNAPSHOT"
    elif pfx in ("TAIL", "SHORT"):
        sk = "TAIL_BACKTEST_USE_KIWOOM_BODY_SNAPSHOT"
    else:
        sk = ""
    if sk and str(get_env_from_db(sk, "")).strip():
        return get_env_bool(sk, False)
    return get_env_bool("BACKTEST_USE_KIWOOM_BODY_SNAPSHOT", False)


def backtest_needs_trigger_snapshot_load(
    params: Optional[Dict[str, Any]] = None,
    *,
    strategy: str = "",
) -> bool:
    """호가/프로그램 스냅샷 DB 로딩이 필요한지 — 필터·본체 재계산 모두 OFF면 스킵.

    필터 OFF 백테에서 ws_orderbook/ws_program/log_backfill 을 통째로 읽으면
    시간만 잡아먹고 손익에는 영향 없다(필터 함수가 조기 return).
    """
    if params is not None and params.get("backtest_use_trigger_snapshot_db") is False:
        return False
    if backtest_use_kiwoom_body_snapshot(params, strategy=strategy):
        return True
    from kis_trader.engine.orderbook_filter import _orderbook_filter_enabled_for_entry
    from kis_trader.engine.program_filter import _program_filter_enabled_for_entry

    p = params or {}
    ob_on = _orderbook_filter_enabled_for_entry(p, strategy)
    pg_on = _program_filter_enabled_for_entry(p, strategy)
    return bool(ob_on or pg_on)


def backtest_use_trigger_snapshot_db(
    params: Optional[Dict[str, Any]] = None,
    *,
    strategy: str = "",
) -> bool:
    """백테·파람서치에서 DB 스냅샷 재생 여부 (기본 true, 필터 OFF면 자동 false)."""
    if params is not None and params.get("backtest_use_trigger_snapshot_db") is not None:
        return bool(params.get("backtest_use_trigger_snapshot_db"))
    if not backtest_needs_trigger_snapshot_load(params, strategy=strategy):
        return False
    pfx = (strategy or "").strip().upper()
    if pfx == "MOMENTUM":
        sk = "MOMENTUM_BACKTEST_USE_TRIGGER_SNAPSHOT_DB"
    elif pfx == "BREAKOUT":
        sk = "BREAKOUT_BACKTEST_USE_TRIGGER_SNAPSHOT_DB"
    elif pfx in ("TAIL", "SHORT"):
        sk = "TAIL_BACKTEST_USE_TRIGGER_SNAPSHOT_DB"
    else:
        sk = ""
    if sk and str(get_env_from_db(sk, "")).strip():
        return get_env_bool(sk, True)
    return get_env_bool("BACKTEST_USE_TRIGGER_SNAPSHOT_DB", True)


def backtest_use_trigger_eval_snapshot(
    params: Optional[Dict[str, Any]] = None,
    *,
    strategy: str = "",
) -> bool:
    """백테에서 filter_eval(판정 시점) 스냅만 사용 (기본 true)."""
    if params is not None and params.get("backtest_use_trigger_eval_snapshot") is not None:
        return bool(params.get("backtest_use_trigger_eval_snapshot"))
    pfx = (strategy or "").strip().upper()
    if pfx == "MOMENTUM":
        sk = "MOMENTUM_BACKTEST_USE_TRIGGER_EVAL_SNAPSHOT"
    elif pfx == "BREAKOUT":
        sk = "BREAKOUT_BACKTEST_USE_TRIGGER_EVAL_SNAPSHOT"
    elif pfx in ("TAIL", "SHORT"):
        sk = "TAIL_BACKTEST_USE_TRIGGER_EVAL_SNAPSHOT"
    else:
        sk = ""
    if sk and str(get_env_from_db(sk, "")).strip():
        return get_env_bool(sk, True)
    return get_env_bool("BACKTEST_USE_TRIGGER_EVAL_SNAPSHOT", True)


def backtest_use_log_backfill_snapshot(
    params: Optional[Dict[str, Any]] = None,
    *,
    strategy: str = "",
) -> bool:
    """백테에서 log_backfill(판정 재생) 사용 여부 (기본 true)."""
    if params is not None and params.get("backtest_use_log_backfill_snapshot") is not None:
        return bool(params.get("backtest_use_log_backfill_snapshot"))
    pfx = (strategy or "").strip().upper()
    if pfx == "MOMENTUM":
        sk = "MOMENTUM_BACKTEST_USE_LOG_BACKFILL_SNAPSHOT"
    elif pfx == "BREAKOUT":
        sk = "BREAKOUT_BACKTEST_USE_LOG_BACKFILL_SNAPSHOT"
    elif pfx in ("TAIL", "SHORT"):
        sk = "TAIL_BACKTEST_USE_LOG_BACKFILL_SNAPSHOT"
    else:
        sk = ""
    if sk and str(get_env_from_db(sk, "")).strip():
        return get_env_bool(sk, True)
    return get_env_bool("BACKTEST_USE_LOG_BACKFILL_SNAPSHOT", True)


def _strategy_db_values(strategy: str) -> Tuple[str, ...]:
    """DB strategy 컬럼 필터 (꼬리=SHORT 로그 별칭 포함)."""
    pfx = (strategy or "").strip().upper()
    if pfx == "MOMENTUM":
        return ("MOMENTUM",)
    if pfx == "BREAKOUT":
        return ("BREAKOUT",)
    if pfx in ("TAIL", "SHORT"):
        return ("TAIL", "SHORT")
    return tuple()


def _normalize_pick_time(bar_candle_time: str, prefer_time: str = "") -> str:
    raw = (prefer_time or bar_candle_time or "").strip()
    digits = raw.replace(":", "").replace("-", "").replace(" ", "")
    if len(digits) >= 14:
        return digits[:14]
    if len(digits) == 12 and digits.isdigit():
        return digits + "59"
    return digits


def _snap_time_key(snap: Any) -> str:
    if isinstance(snap, dict):
        st = snap.get("snap_time") or ""
    else:
        st = getattr(snap, "snap_time", None) or ""
    return str(st).strip()[:14]


def _pick_latest_le(bucket: List[Any], target_time: str) -> Any:
    """target_time 이하 중 가장 최근 스냅샷 (판정 시점 재생)."""
    if not bucket:
        return None
    target = _normalize_pick_time(target_time, target_time)
    best = None
    best_key = ""
    for snap in bucket:
        sk = _snap_time_key(snap)
        if not sk:
            continue
        if sk <= target and sk >= best_key:
            best = snap
            best_key = sk
    return best or bucket[-1]


def _minute_key(snap_time: str) -> str:
    st = str(snap_time or "").strip()
    return st[:12] if len(st) >= 12 else st


def _load_orderbook_rows(
    db,
    start_key: str,
    end_key: str,
    codes: Optional[Set[str]],
    *,
    market: str,
    sources: Tuple[str, ...],
    strategy: str = "",
) -> List[Dict[str, Any]]:
    """ws_orderbook 행 조회 (source·strategy 필터)."""
    tt_start = (start_key or "")[:14]
    if len(tt_start) == 12:
        tt_start += "00"
    tt_end = (end_key or "")[:14]
    if len(tt_end) == 12:
        tt_end += "59"

    code_filter = ""
    params: List[Any] = [market, tt_start, tt_end]
    if codes:
        placeholders = ",".join(["%s"] * len(codes))
        code_filter = f" AND code IN ({placeholders})"
        params.extend(sorted(codes))

    source_filter = ""
    if sources:
        source_filter = f" AND source IN ({','.join(['%s'] * len(sources))})"
        params.extend(sources)

    strat_vals = _strategy_db_values(strategy)
    strat_filter = ""
    if strat_vals:
        strat_filter = (
            f" AND (strategy IS NULL OR strategy IN ({','.join(['%s'] * len(strat_vals))}))"
        )
        params.extend(strat_vals)

    sql = f"""
        SELECT code, snap_time, best_bid, best_ask, total_bid_qty, total_ask_qty,
               bid_qty_l3, ask_qty_l3, levels_json, source, recv_ts,
               strategy, reject_code, reject_msg
        FROM ws_orderbook
        WHERE market = %s
          AND snap_time >= %s
          AND snap_time <= %s
          {code_filter}
          {source_filter}
          {strat_filter}
        ORDER BY code, snap_time
    """
    try:
        rows = db.conn.execute(sql, tuple(params)).fetchall()
        return [dict(r) for r in rows]
    except Exception as e:
        logger.warning("ws_orderbook 조회 실패 — TRIGGER 호가 스킵: %s", e)
        return []


def _index_log_verdict_row(
    out: Dict[str, Dict[str, List[Dict[str, Any]]]],
    row: Dict[str, Any],
) -> None:
    """log_backfill 판정 1건 → code/minute 버킷."""
    code = str(row.get("code") or "").strip()
    st = str(row.get("snap_time") or "").strip()
    minute_key = _minute_key(st)
    if not code or not minute_key:
        return
    bucket = out[code].setdefault(minute_key, [])
    bucket.append({
        "snap_time": st[:14],
        "reject_code": row.get("reject_code"),
        "reject_msg": row.get("reject_msg") or "",
        "strategy": row.get("strategy"),
    })


def load_orderbook_snapshots_by_code(
    db,
    start_key: str,
    end_key: str,
    codes: Optional[Set[str]] = None,
    *,
    market: Optional[str] = None,
    engine_params: Optional[Dict[str, Any]] = None,
    strategy: str = "",
) -> Tuple[Dict[str, Dict[str, List[OrderbookSnapshot]]], int, Dict[str, Dict[str, List[Dict[str, Any]]]]]:
    """기간 내 ``ws_orderbook`` → ``{code: {minute: [snap, ...]}}`` + log_backfill 판정."""
    mkt = (market or get_env_from_db("WS_TICK_DEFAULT_MARKET", "KR") or "KR").strip().upper()
    try:
        if hasattr(db, "ensure_ws_orderbook_table"):
            db.ensure_ws_orderbook_table()
    except Exception:
        pass

    use_eval = backtest_use_trigger_eval_snapshot(engine_params, strategy=strategy)
    use_log_bf = backtest_use_log_backfill_snapshot(engine_params, strategy=strategy)
    use_kiwoom = backtest_use_kiwoom_body_snapshot(engine_params, strategy=strategy)

    sources: List[str] = []
    if use_eval:
        sources.append(FILTER_EVAL_SOURCE)
    if use_log_bf:
        sources.append(LOG_BACKFILL_SOURCE)
    if use_kiwoom:
        sources.append(KIWOOM_BODY_SOURCE)
    if not sources:
        return {}, 0, {}

    rows = _load_orderbook_rows(
        db, start_key, end_key, codes, market=mkt,
        sources=tuple(sources), strategy=strategy,
    )

    out: Dict[str, Dict[str, List[OrderbookSnapshot]]] = defaultdict(dict)
    log_verdict: Dict[str, Dict[str, List[Dict[str, Any]]]] = defaultdict(dict)
    seen_eval: Set[Tuple[str, str]] = set()
    # 본체(kiwoom_0d) 다운샘플: (code, minute) 당 최신 1건만 유지 (분단위 백테·메모리 절약)
    kiwoom_latest: Dict[str, Dict[str, str]] = defaultdict(dict)
    total = 0

    for r in rows:
        code = str(r.get("code") or "").strip()
        src = str(r.get("source") or "").strip()
        st = str(r.get("snap_time") or "").strip()
        minute_key = _minute_key(st)
        if not code or not minute_key:
            continue

        if src == FILTER_EVAL_SOURCE:
            snap = orderbook_snapshot_from_storage(r)
            bucket = out[code].setdefault(minute_key, [])
            bucket.append(snap)
            seen_eval.add((code, st[:14]))
            total += 1
        elif src == KIWOOM_BODY_SOURCE:
            # 분당 최신 1건만 (rows 는 snap_time ASC 정렬 → 뒤가 최신)
            prev = kiwoom_latest[code].get(minute_key)
            snap = orderbook_snapshot_from_storage(r)
            bucket = out[code].setdefault(minute_key, [])
            if prev is None:
                bucket.append(snap)
                total += 1
            else:
                bucket[-1] = snap  # 같은 분 더 늦은 스냅으로 교체
            kiwoom_latest[code][minute_key] = st[:14]
        elif src == LOG_BACKFILL_SOURCE:
            _index_log_verdict_row(log_verdict, r)
            # 본체 재계산 모드면 log_backfill 빈본체는 out 에 넣지 않는다
            # (kiwoom 본체가 body source 가 되도록). 판정(verdict)은 폴백용으로 보존.
            if use_kiwoom:
                continue
            if (code, st[:14]) in seen_eval:
                continue
            snap = orderbook_snapshot_from_storage(r)
            bucket = out[code].setdefault(minute_key, [])
            bucket.append(snap)
            total += 1

    return dict(out), total, dict(log_verdict)


def load_program_snapshots_by_code(
    db,
    start_key: str,
    end_key: str,
    codes: Optional[Set[str]] = None,
    *,
    market: Optional[str] = None,
    engine_params: Optional[Dict[str, Any]] = None,
    strategy: str = "",
) -> Tuple[Dict[str, Dict[str, List[ProgramSnapshot]]], int]:
    """기간 내 ``ws_program`` → ``{code: {minute: [snap, ...]}}``."""
    mkt = (market or get_env_from_db("WS_TICK_DEFAULT_MARKET", "KR") or "KR").strip().upper()
    tt_start = (start_key or "")[:14]
    if len(tt_start) == 12:
        tt_start += "00"
    tt_end = (end_key or "")[:14]
    if len(tt_end) == 12:
        tt_end += "59"
    try:
        if hasattr(db, "ensure_ws_program_table"):
            db.ensure_ws_program_table()
    except Exception:
        pass

    code_filter = ""
    params: List[Any] = [mkt, tt_start, tt_end]
    if codes:
        placeholders = ",".join(["%s"] * len(codes))
        code_filter = f" AND code IN ({placeholders})"
        params.extend(sorted(codes))

    source_filter = ""
    if backtest_use_trigger_eval_snapshot(engine_params, strategy=strategy):
        source_filter = " AND source = %s"
        params.append(FILTER_EVAL_SOURCE)

    sql = f"""
        SELECT code, snap_time, buy_qty, sell_qty, net_qty,
               buy_amt, sell_amt, net_amt, source, recv_ts
        FROM ws_program
        WHERE market = %s
          AND snap_time >= %s
          AND snap_time <= %s
          {code_filter}
          {source_filter}
        ORDER BY code, snap_time
    """
    try:
        rows = db.conn.execute(sql, tuple(params)).fetchall()
    except Exception as e:
        logger.warning("ws_program 조회 실패 — TRIGGER 프로그램 스킵: %s", e)
        return {}, 0

    out: Dict[str, Dict[str, List[ProgramSnapshot]]] = defaultdict(dict)
    total = 0
    for r in rows:
        code = str(r["code"]).strip()
        minute_key = _minute_key(str(r["snap_time"]))
        if not minute_key:
            continue
        snap = program_snapshot_from_storage(dict(r))
        bucket = out[code].setdefault(minute_key, [])
        bucket.append(snap)
        total += 1
    return dict(out), total


def snapshot_coverage_stats(
    orderbook_by_code: Dict[str, Dict[str, List[OrderbookSnapshot]]],
    program_by_code: Dict[str, Dict[str, List[ProgramSnapshot]]],
) -> Dict[str, Any]:
    ob_codes = len(orderbook_by_code)
    pg_codes = len(program_by_code)
    ob_minutes = sum(len(v) for v in orderbook_by_code.values())
    pg_minutes = sum(len(v) for v in program_by_code.values())
    return {
        "orderbook_codes_with_data": ob_codes,
        "program_codes_with_data": pg_codes,
        "orderbook_minute_buckets": ob_minutes,
        "program_minute_buckets": pg_minutes,
    }


def load_trigger_snapshots_by_code(
    db,
    start_key: str,
    end_key: str,
    codes: Optional[Set[str]] = None,
    *,
    market: Optional[str] = None,
    engine_params: Optional[Dict[str, Any]] = None,
    strategy: str = "",
) -> Tuple[
    Dict[str, Dict[str, List[OrderbookSnapshot]]],
    Dict[str, Dict[str, List[ProgramSnapshot]]],
    Dict[str, Any],
]:
    """호가·프로그램 스냅샷 일괄 로드 + 커버리지 메타."""
    ob_by_code, ob_rows, log_verdict = load_orderbook_snapshots_by_code(
        db, start_key, end_key, codes, market=market,
        engine_params=engine_params, strategy=strategy,
    )
    pg_by_code, pg_rows = load_program_snapshots_by_code(
        db, start_key, end_key, codes, market=market,
        engine_params=engine_params, strategy=strategy,
    )
    meta = snapshot_coverage_stats(ob_by_code, pg_by_code)
    meta["ws_orderbook_rows_loaded"] = ob_rows
    meta["ws_program_rows_loaded"] = pg_rows
    meta["trigger_eval_only"] = backtest_use_trigger_eval_snapshot(
        engine_params, strategy=strategy,
    )
    meta["log_backfill_enabled"] = backtest_use_log_backfill_snapshot(
        engine_params, strategy=strategy,
    )
    meta["log_verdict_by_code"] = log_verdict
    meta["log_verdict_rows"] = sum(
        len(lst) for cm in log_verdict.values() for lst in cm.values()
    )
    return ob_by_code, pg_by_code, meta


def pick_snapshot_for_bar(
    by_code: Optional[Dict[str, Dict[str, Any]]],
    code: str,
    bar_candle_time: str,
    *,
    prefer_time: str = "",
) -> Any:
    """신호/진입 시각에 맞는 최신 스냅샷 1건 (판정 시점 ≤ prefer_time)."""
    if not by_code:
        return None
    target = _normalize_pick_time(bar_candle_time, prefer_time)
    minute_key = target[:12]
    bucket = (by_code.get(code) or {}).get(minute_key) or []
    if bucket:
        return _pick_latest_le(bucket, target)
    code_map = by_code.get(code) or {}
    prior = [k for k in code_map.keys() if k <= minute_key]
    if not prior:
        return None
    prev_key = max(prior)
    prev_bucket = code_map.get(prev_key) or []
    return _pick_latest_le(prev_bucket, target) if prev_bucket else None


def pick_verdict_for_bar(
    by_code: Optional[Dict[str, Dict[str, List[Dict[str, Any]]]]],
    code: str,
    bar_candle_time: str,
    *,
    prefer_time: str = "",
) -> Optional[Dict[str, Any]]:
    """log_backfill 판정 1건 (판정 시점 ≤ prefer_time)."""
    if not by_code:
        return None
    target = _normalize_pick_time(bar_candle_time, prefer_time)
    minute_key = target[:12]
    bucket = (by_code.get(code) or {}).get(minute_key) or []
    if bucket:
        best = _pick_latest_le(bucket, target)
        return best if isinstance(best, dict) else None
    code_map = by_code.get(code) or {}
    prior = [k for k in code_map.keys() if k <= minute_key]
    if not prior:
        return None
    prev_key = max(prior)
    prev_bucket = code_map.get(prev_key) or []
    best = _pick_latest_le(prev_bucket, target)
    return best if isinstance(best, dict) else None


def inject_trigger_snapshots_into_params(
    params: Dict[str, Any],
    *,
    orderbook_by_code: Optional[Dict[str, Dict[str, List[OrderbookSnapshot]]]] = None,
    program_by_code: Optional[Dict[str, Dict[str, List[ProgramSnapshot]]]] = None,
    log_verdict_by_code: Optional[Dict[str, Dict[str, List[Dict[str, Any]]]]] = None,
    code: str,
    bar_candle_time: str,
    prefer_time: str = "",
) -> None:
    """백테 포트폴리오 — eval 직전 params 에 호가·프로그램 스냅샷·log 판정 주입."""
    pick_t = prefer_time or bar_candle_time
    ob = pick_snapshot_for_bar(orderbook_by_code, code, bar_candle_time, prefer_time=pick_t)
    if ob is not None:
        params["_backtest_orderbook_snapshot"] = ob
    else:
        params.pop("_backtest_orderbook_snapshot", None)
    pg = pick_snapshot_for_bar(program_by_code, code, bar_candle_time, prefer_time=pick_t)
    if pg is not None:
        params["_backtest_program_snapshot"] = pg
    else:
        params.pop("_backtest_program_snapshot", None)

    lv_src = log_verdict_by_code
    if lv_src is None:
        lv_src = params.get("_backtest_log_verdict_by_code")
    verdict = pick_verdict_for_bar(lv_src, code, bar_candle_time, prefer_time=pick_t)
    if verdict is not None:
        rc = verdict.get("reject_code")
        rm = str(verdict.get("reject_msg") or "")
        params["_backtest_log_orderbook_verdict"] = (rc, rm) if rc else (None, None)
    else:
        params.pop("_backtest_log_orderbook_verdict", None)