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
    
    chunks = list(_iter_day_chunks_dt(s_dt, e_dt))
    if not chunks:
        return {}, 0
        
    max_w = min(10, len(chunks)) if chunks else 1
    max_w = int(max(1, max_w))
    from concurrent.futures import ThreadPoolExecutor, as_completed

    def process_chunk(c_s, c_e):
        from database import TradeDB
        local_db = TradeDB()
        local_out = defaultdict(dict)
        n = 0
        try:
            if code_list:
                ph = ",".join(["%s"] * len(code_list))
                rows = local_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})",
                    [c_s, c_e, *code_list],
                ).fetchall()
            else:
                rows = local_db.conn.execute(
                    "SELECT code, ts, price, volume, chetime FROM ls_ws_ticks "
                    "WHERE ts >= %s AND ts <= %s",
                    [c_s, c_e],
                ).fetchall()
                
            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())
                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 = local_out[code].setdefault(minute_key, [])
                bucket.append(tick)
                n += 1
        except Exception as e:
            logger.warning("ls_ws_ticks day=%s 조회 실패: %s", c_s[:10], e)
        finally:
            local_db.close()
        return c_s, c_e, local_out, n

    with ThreadPoolExecutor(max_workers=max_w) as executor:
        futures = {executor.submit(process_chunk, c[0], c[1]): c for c in chunks}
        for future in as_completed(futures):
            c_s, c_e = futures[future]
            try:
                _, _, local_out, n = future.result()
                for _code, mins_dict in local_out.items():
                    code_bucket = out[_code]
                    for min_k, ticks in mins_dict.items():
                        code_bucket.setdefault(min_k, []).extend(ticks)
                total += n
            except Exception as e:
                logger.warning("ls_ws_ticks day=%s 스레드 실패: %s", c_s[:10], e)

    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