kis_bot/kis_trader/backtest/ls_history_loaders.py

"""
kis_trader/backtest/ls_history_loaders.py — ls_ws_candles / ls_ws_ticks 백테 로드
==============================================================================
``history_source=ls`` 일 때 키움 ``ws_*`` 대신 LS 테이블을 읽는다.
반환 스키마는 전략/엔진이 기대하는 candle_time(12)·tick_time(14) 으로 정규화.
"""
from __future__ import annotations

from collections import defaultdict
from datetime import datetime, timedelta
from typing import Any, Dict, Iterator, List, Optional, Set, Tuple

from kis_trader.utils.logger import get_logger

logger = get_logger("kis_trader.ls_history_loaders")


def _ct_to_ls_dt(ct: str) -> str:
    d = "".join(ch for ch in (ct or "") if ch.isdigit())[:12]
    if len(d) < 12:
        return ""
    return f"{d[0:4]}-{d[4:6]}-{d[6:8]} {d[8:10]}:{d[10:12]}:00"


def _ls_dt_to_ct(dt_s: str) -> str:
    s = (dt_s or "").strip()
    if len(s) >= 16 and s[4] == "-" and s[10] == " ":
        return s[0:4] + s[5:7] + s[8:10] + s[11:13] + s[14:16]
    digits = "".join(ch for ch in s if ch.isdigit())
    return digits[:12]


def _row_to_strategy_candle(r: dict) -> dict:
    ct = _ls_dt_to_ct(str(r.get("datetime") or ""))
    return {
        "candle_time": ct,
        "open": float(r.get("open") or 0),
        "high": float(r.get("high") or 0),
        "low": float(r.get("low") or 0),
        "close": float(r.get("close") or 0),
        "volume": float(r.get("volume") or 0),
        "is_confirmed": 1,
        "source": "ls",
    }


def load_ls_candles_by_code(
    db,
    start_key: str,
    end_key: str,
    *,
    min_bars: int = 5,
    tf_min: int = 1,
) -> Tuple[Dict[str, List[Dict]], int]:
    """ls_ws_candles → {code: [candle dict...]}, total_rows."""
    s_dt = _ct_to_ls_dt(start_key)
    e_dt = _ct_to_ls_dt(end_key)
    if not s_dt or not e_dt:
        return {}, 0
    tf = max(1, int(tf_min or 1))
    mb = max(1, int(min_bars or 5))
    codes_raw = db.conn.execute(
        "SELECT DISTINCT code FROM ls_ws_candles WHERE tf_min=%s "
        "AND datetime >= %s AND datetime <= %s ORDER BY code",
        [tf, s_dt, e_dt],
    ).fetchall()
    codes = [r["code"] for r in (codes_raw or [])]
    candles_by_code: Dict[str, List[Dict]] = {}
    total = 0
    for code in codes:
        rows = db.conn.execute(
            "SELECT datetime, open, high, low, close, volume "
            "FROM ls_ws_candles WHERE tf_min=%s AND code=%s "
            "AND datetime >= %s AND datetime <= %s "
            "ORDER BY datetime ASC",
            [tf, code, s_dt, e_dt],
        ).fetchall()
        if len(rows or []) < mb:
            continue
        bars = [_row_to_strategy_candle(dict(r)) for r in rows]
        bars = [b for b in bars if len(str(b.get("candle_time") or "")) >= 12]
        if len(bars) < mb:
            continue
        candles_by_code[code] = bars
        total += len(bars)
    return candles_by_code, total


def prepend_ls_candle_warmup(
    db,
    candles_by_code: Dict[str, List[Dict]],
    period_start_key: str,
    warmup_bars: int,
    *,
    tf_min: int = 1,
) -> int:
    """기간 시작 전 N봉 prepend (RSI 등)."""
    wb = max(0, int(warmup_bars or 0))
    if wb <= 0 or not period_start_key or not candles_by_code:
        return 0
    ps = str(period_start_key)[:12]
    ps_dt = _ct_to_ls_dt(ps)
    if not ps_dt:
        return 0
    tf = max(1, int(tf_min or 1))
    total_prepended = 0
    for code, rows in list(candles_by_code.items()):
        if not rows:
            continue
        first_period_idx = None
        for i, r in enumerate(rows):
            ct = str(r.get("candle_time") or "")
            if ct >= ps:
                first_period_idx = i
                break
        if first_period_idx is None or first_period_idx > 0:
            continue
        first_ct = str(rows[first_period_idx].get("candle_time") or "")
        if not first_ct:
            continue
        first_dt = _ct_to_ls_dt(first_ct)
        warm_rows = db.conn.execute(
            "SELECT datetime, open, high, low, close, volume "
            "FROM ls_ws_candles WHERE tf_min=%s AND code=%s "
            "AND datetime < %s ORDER BY datetime DESC LIMIT %s",
            [tf, code, first_dt, wb],
        ).fetchall()
        if not warm_rows:
            continue
        prefix = [_row_to_strategy_candle(dict(r)) for r in reversed(list(warm_rows))]
        candles_by_code[code] = prefix + [dict(r) for r in rows]
        total_prepended += len(prefix)
    return total_prepended


def _iter_day_chunks_dt(s_dt: str, e_dt: str) -> Iterator[Tuple[str, str]]:
    """datetime 문자열 범위를 달력일 단위로."""
    try:
        d0 = datetime.strptime(s_dt[:10], "%Y-%m-%d")
        d1 = datetime.strptime(e_dt[:10], "%Y-%m-%d")
    except ValueError:
        return
    cur = d0
    while cur <= d1:
        day = cur.strftime("%Y-%m-%d")
        chunk_s = max(s_dt, f"{day} 00:00:00")
        chunk_e = min(e_dt, f"{day} 23:59:59")
        if chunk_s <= chunk_e:
            yield chunk_s, chunk_e
        cur += timedelta(days=1)


def _ts_to_tick_time(ts) -> str:
    if isinstance(ts, datetime):
        return ts.strftime("%Y%m%d%H%M%S")
    s = str(ts or "").strip()
    digits = "".join(ch for ch in s if ch.isdigit())
    if len(digits) >= 14:
        return digits[:14]
    # "YYYY-MM-DD HH:MM:SS"
    if len(s) >= 19 and s[4] == "-":
        return (
            s[0:4] + s[5:7] + s[8:10]
            + s[11:13] + s[14:16] + s[17:19]
        )
    return digits[:14].ljust(14, "0")


def load_ls_ticks_by_code(
    db,
    start_key: str,
    end_key: str,
    codes: Optional[Set[str]] = None,
) -> Tuple[Dict[str, Dict[str, List[Dict[str, Any]]]], int]:
    """
    ls_ws_ticks → {code: {minute_key: [tick...]}}.
    tick 스키마: tick_time/price/volume/source(=ls)/_lag_sec(가능 시)
    """
    from kis_trader.engine.feed_fallback import packet_lag_seconds

    s12 = (start_key or "")[:12]
    e12 = (end_key or "")[:12]
    s_dt = _ct_to_ls_dt(s12)
    e_dt = _ct_to_ls_dt(e12)
    if not s_dt or not e_dt:
        return {}, 0
    # end 분 포함 → 초 59
    e_dt = e_dt[:17] + "59" if len(e_dt) >= 17 else e_dt
    out: Dict[str, Dict[str, List[Dict[str, Any]]]] = defaultdict(dict)
    total = 0
    code_list = sorted(codes) if codes else None
    for chunk_s, chunk_e in _iter_day_chunks_dt(s_dt, e_dt):
        try:
            if code_list:
                # 종목 수가 많으면 IN 절 — 청크당 한 번
                ph = ",".join(["%s"] * len(code_list))
                rows = db.conn.execute(
                    f"SELECT code, ts, price, volume, chetime FROM ls_ws_ticks "
                    f"WHERE ts >= %s AND ts <= %s AND code IN ({ph})",
                    [chunk_s, chunk_e, *code_list],
                ).fetchall()
            else:
                rows = db.conn.execute(
                    "SELECT code, ts, price, volume, chetime FROM ls_ws_ticks "
                    "WHERE ts >= %s AND ts <= %s",
                    [chunk_s, chunk_e],
                ).fetchall()
        except Exception as e:
            logger.warning("ls_ws_ticks day=%s 조회 실패: %s", chunk_s[:10], e)
            continue
        for r in rows or []:
            code = str(r["code"]).strip()
            tt = _ts_to_tick_time(r.get("ts"))
            che = str(r.get("chetime") or "").strip()
            che_d = "".join(ch for ch in che if ch.isdigit())
            # chetime 이 HHMMSS 이면 수신 ts 날짜와 합쳐 14자리
            if len(che_d) >= 14:
                tt = che_d[:14]
            elif len(che_d) >= 6 and len(tt) >= 8:
                tt = (tt[:8] + che_d[-6:]).ljust(14, "0")[:14]
            if len(tt) < 12:
                continue
            minute_key = tt[:12]
            lag = None
            try:
                recv_raw = r.get("ts")
                if isinstance(recv_raw, datetime):
                    recv_dt = recv_raw
                else:
                    recv_s = str(recv_raw or "")[:19]
                    recv_dt = datetime.strptime(recv_s, "%Y-%m-%d %H:%M:%S")
                lag_f = packet_lag_seconds(tt, now_dt=recv_dt)
                if lag_f is not None:
                    lag = int(lag_f)
            except Exception:
                lag = None
            tick = {
                "code": code,
                "tick_time": tt.ljust(14, "0")[:14],
                "price": float(r.get("price") or 0),
                "volume": int(float(r.get("volume") or 0)),
                "source": "ls",
                "_lag_sec": lag,
            }
            bucket = out[code].setdefault(minute_key, [])
            bucket.append(tick)
            total += 1
    for _code, minutes in out.items():
        for _mk, ticks in minutes.items():
            ticks.sort(key=lambda t: str(t.get("tick_time") or ""))
    return dict(out), total