kis_bot/kis_trader/backtest/bt_candle_source.py

#!/usr/bin/env python3
"""
백테·Optuna ws_candles 소스 선택 — 실매 get_candles 와 동일.

읽기: 틱 메인 WS 한 칸, 없으면 kiwoom+rest 구멍, 그다음 kiwoom+rollup.
CANDLE_SOURCE=kis|kiwoom 이면 그 메인(+구멍). 빈값이면 LIVE_TICK_PROVIDER.
"""
from __future__ import annotations

from typing import Any, Dict, List, Optional, Sequence, Tuple

from kis_trader.ws.candle_series import (
    dedupe_by_read_pairs,
    live_read_label,
    live_read_pairs,
    normalize_source_channel,
)

ReadPair = Tuple[str, str]

# 레거시 호환 이름 (실제 필터는 live_read_pairs)
BT_WS_CANDLE_SOURCES: Tuple[str, ...] = ("kiwoom", "kis", "rest", "rollup_1m")


def resolve_bt_candle_source_override() -> str:
    """'' = LIVE 병합, 'kis'|'kiwoom' = 그 메인(+키움 REST 구멍)."""
    from kis_trader.ws.candle_series import _provider_and_override
    _provider, override = _provider_and_override()
    return override


def live_candle_source_order() -> Tuple[str, ...]:
    """호환: source 문자열만. 실제 조회는 resolve_bt_read_pairs."""
    return tuple(dict.fromkeys(p[0] for p in live_read_pairs()))


def resolve_bt_read_pairs() -> Tuple[ReadPair, ...]:
    return live_read_pairs()


def resolve_bt_candle_source_order() -> Tuple[str, ...]:
    """호환 래퍼 — 신규 코드는 resolve_bt_read_pairs 사용."""
    return live_candle_source_order()


def resolve_bt_candle_source_label() -> str:
    return live_read_label()


def dedupe_candle_rows(
    rows: Sequence[Dict[str, Any]],
    source_order: Optional[Sequence[Any]] = None,
) -> List[Dict[str, Any]]:
    """candle_time 1행. source_order 가 쌍이면 그대로, 아니면 live_read_pairs."""
    pairs: Optional[Sequence[ReadPair]] = None
    if source_order:
        first = source_order[0]
        if isinstance(first, (tuple, list)) and len(first) >= 2:
            pairs = tuple((str(a), str(b)) for a, b in source_order)  # type: ignore[misc]
        else:
            # 레거시 source 문자열 → 정규화 후 순위
            mapped: List[ReadPair] = []
            for s in source_order:
                mapped.append(normalize_source_channel(str(s), ""))
            pairs = tuple(mapped)
    return dedupe_by_read_pairs(rows, pairs)


def _pairs_in_sql(pairs: Sequence[ReadPair]) -> Tuple[str, List[str]]:
    parts: List[str] = []
    params: List[str] = []
    for src, ch in pairs:
        parts.append("(source=%s AND channel=%s)")
        params.extend([src, ch])
    return " AND (" + " OR ".join(parts) + ")", params


def _ticks_for_bar_garbage(db, code: str, start_key: str, end_key: str) -> Optional[List[Dict[str, Any]]]:
    from kis_trader.engine.feed_fallback import candle_garbage_fallback_enabled

    if not candle_garbage_fallback_enabled():
        return None
    sk = str(start_key or "").strip()
    ek = str(end_key or "").strip()
    if len(sk) == 12:
        sk = sk + "00"
    if len(ek) == 12:
        ek = ek + "59"
    try:
        rows = db.conn.execute(
            "SELECT tick_time, source, tick_time_raw FROM ws_ticks "
            "WHERE code=%s AND tick_time >= %s AND tick_time <= %s",
            (str(code).strip(), sk, ek),
        ).fetchall()
        return [dict(r) for r in (rows or [])]
    except Exception:
        return None


def list_ws_candle_codes(
    db,
    timeframe: int,
    start_key: str,
    end_key: str,
    *,
    market: Optional[str] = None,
) -> List[str]:
    """기간 내 종목 코드 — 선택된 (source,channel) 기준 DISTINCT."""
    pairs = resolve_bt_read_pairs()
    src_sql, src_params = _pairs_in_sql(pairs)
    mk = (market or "").strip().upper()
    if mk:
        rows = db.conn.execute(
            "SELECT DISTINCT code FROM ws_candles WHERE timeframe=%s AND market=%s "
            "AND candle_time >= %s AND candle_time <= %s"
            + src_sql
            + " ORDER BY code",
            [int(timeframe), mk, start_key, end_key, *src_params],
        ).fetchall()
    else:
        rows = db.conn.execute(
            "SELECT DISTINCT code FROM ws_candles WHERE timeframe=%s "
            "AND candle_time >= %s AND candle_time <= %s"
            + src_sql
            + " ORDER BY code",
            [int(timeframe), start_key, end_key, *src_params],
        ).fetchall()
    return [r["code"] for r in rows]


def fetch_ws_candles_for_code(
    db,
    code: str,
    timeframe: int,
    start_key: str,
    end_key: str,
    *,
    extra_select: str = "",
    peak_sel: str = "",
    market: Optional[str] = None,
    confirmed_only: bool = True,
) -> List[Dict[str, Any]]:
    """단일 종목·기간 봉 로드 — 메인 WS 우선, 구멍 kiwoom+rest(+rollup)."""
    pairs = resolve_bt_read_pairs()
    confirmed_sql = " AND is_confirmed=1" if confirmed_only else ""
    mk = (market or "").strip().upper()
    src_sql, src_params = _pairs_in_sql(pairs)
    cols = (
        f"candle_time, open, high, low, close, volume, source, channel"
        f"{peak_sel}{extra_select}"
    )
    if mk:
        rows = db.conn.execute(
            f"SELECT {cols} FROM ws_candles "
            "WHERE timeframe=%s AND code=%s AND market=%s "
            "AND candle_time >= %s AND candle_time <= %s"
            + confirmed_sql
            + src_sql
            + " ORDER BY candle_time ASC",
            [int(timeframe), code, mk, start_key, end_key, *src_params],
        ).fetchall()
    else:
        rows = db.conn.execute(
            f"SELECT {cols} FROM ws_candles "
            "WHERE timeframe=%s AND code=%s "
            "AND candle_time >= %s AND candle_time <= %s"
            + confirmed_sql
            + src_sql
            + " ORDER BY candle_time ASC",
            [int(timeframe), code, start_key, end_key, *src_params],
        ).fetchall()
    ticks = _ticks_for_bar_garbage(db, str(code), start_key, end_key)
    return dedupe_by_read_pairs(
        [dict(r) for r in rows], pairs, ticks=ticks, tf_min=int(timeframe),
        missing_policy="hole", live_cover=False,
    )


def fetch_ws_candles_warmup_before(
    db,
    code: str,
    timeframe: int,
    before_candle_time: str,
    limit: int,
    *,
    extra_select: str = "",
    peak_sel: str = "",
    confirmed_only: bool = True,
) -> List[Dict[str, Any]]:
    """기간 시작 이전 N봉 — prepend 웜업용 (오래된→최신)."""
    if limit <= 0 or not before_candle_time:
        return []
    pairs = resolve_bt_read_pairs()
    confirmed_sql = " AND is_confirmed=1" if confirmed_only else ""
    fetch_limit = max(int(limit) * max(len(pairs), 1), int(limit) + 50)
    src_sql, src_params = _pairs_in_sql(pairs)
    cols = (
        f"candle_time, open, high, low, close, volume, source, channel"
        f"{peak_sel}{extra_select}"
    )
    rows = db.conn.execute(
        f"SELECT {cols} FROM ws_candles "
        "WHERE timeframe=%s AND code=%s AND candle_time < %s"
        + confirmed_sql
        + src_sql
        + " ORDER BY candle_time DESC LIMIT %s",
        [int(timeframe), code, before_candle_time, *src_params, fetch_limit],
    ).fetchall()
    raw_rows = [dict(r) for r in reversed(rows)]
    t0 = str(raw_rows[0].get("candle_time") or "")[:12] if raw_rows else ""
    ticks = (
        _ticks_for_bar_garbage(db, str(code), t0, str(before_candle_time))
        if t0 else None
    )
    bars = dedupe_by_read_pairs(
        raw_rows, pairs,
        ticks=ticks,
        tf_min=int(timeframe),
        missing_policy="hole",
        live_cover=False,
    )
    return bars[-limit:] if len(bars) > limit else bars