kis_bot/kis_trader/ws/candle_series.py

"""
ws_candles 시리즈 — source=증권사, channel=ws|rest|rollup.

저장은 섞지 않음: 키움 REST 갭보정은 항상 source=kiwoom, channel=rest.
읽기(실매=Optuna): 메인 WS → 2차 WS → LS WS → kiwoom+rest → rollup.
"""
from __future__ import annotations

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

from kis_trader.utils.env import get_env_from_db

CHANNEL_WS = "ws"
CHANNEL_REST = "rest"
CHANNEL_ROLLUP = "rollup"

# 읽기 우선순위 쌍 (source, channel)
ReadPair = Tuple[str, str]


def normalize_source_channel(
    source: str = "",
    channel: str = "",
) -> Tuple[str, str]:
    """레거시 source=rest|kw_rest|rollup_1m 을 (증권사, channel)로 정규화."""
    s = str(source or "").strip().lower()
    ch = str(channel or "").strip().lower()
    if s in ("rest", "kw_rest"):
        return "kiwoom", CHANNEL_REST
    if s in ("rollup_1m", "rollup"):
        src = "kiwoom"
        if ch in ("kis",):
            src = "kis"
        return src, CHANNEL_ROLLUP
    if s in ("kis", "kis_us", "ws", ""):
        return "kis", (ch or CHANNEL_WS)
    if s == "kiwoom":
        return "kiwoom", (ch or CHANNEL_WS)
    if s == "ls":
        return "ls", (ch or CHANNEL_WS)
    return (s or "kis")[:16], (ch or CHANNEL_WS)[:10]


def _provider_and_override() -> Tuple[str, str]:
    override = ""
    raw = os.environ.get("CANDLE_SOURCE")
    if raw is None:
        try:
            raw = get_env_from_db("CANDLE_SOURCE", "")
        except Exception:
            raw = ""
    override = str(raw or "").strip().lower()
    if override not in ("kis", "kiwoom"):
        override = ""
    try:
        provider = (
            get_env_from_db("LIVE_TICK_PROVIDER", "kiwoom") or "kiwoom"
        ).strip().lower()
    except Exception:
        provider = "kiwoom"
    if provider not in ("kis", "kiwoom"):
        provider = "kiwoom"
    return provider, override


def live_read_pairs() -> Tuple[ReadPair, ...]:
    """실매·Optuna 기본 읽기 순서.

    메인 WS → 2차 WS → LS WS(있으면) → 키움 REST 구멍 → 롤업.
    CANDLE_SOURCE=kis|kiwoom 이면 그 메인을 씀. 빈값이면 LIVE_TICK_PROVIDER.
    """
    provider, override = _provider_and_override()
    main = override or provider
    aux = "kis" if main == "kiwoom" else "kiwoom"
    return (
        (main, CHANNEL_WS),
        (aux, CHANNEL_WS),
        ("ls", CHANNEL_WS),
        ("kiwoom", CHANNEL_REST),
        ("kiwoom", CHANNEL_ROLLUP),
    )


def live_read_label() -> str:
    provider, override = _provider_and_override()
    main = override or provider
    aux = "kis" if main == "kiwoom" else "kiwoom"
    if override:
        return f"{override.upper()}+{aux}+ls+kiwoom_rest"
    return f"LIVE({main})+{aux}+ls+kiwoom_rest"


def dedupe_by_read_pairs(
    rows: Sequence[Dict[str, Any]],
    pairs: Optional[Sequence[ReadPair]] = None,
    *,
    ticks: Optional[Sequence[Dict[str, Any]]] = None,
    tf_min: int = 1,
    missing_policy: str = "hole",
    live_cover: bool = False,
) -> List[Dict[str, Any]]:
    """candle_time 1행 — pairs 앞쪽이 이김. source/channel 은 결과에 남기지 않음.

    ticks 가 있으면 WS 봉만 쓰레기(0건·봉끝 나이초과)일 때 다음 쌍. 한 봉 혼합 없음.
    live_cover=True 이면 링이 그 분을 커버 못할 때 그 WS 봉을 유지.
    """
    from kis_trader.engine.feed_fallback import bar_is_garbage, live_bar_is_garbage

    order = tuple(pairs or live_read_pairs())
    by_ct: Dict[str, Dict[ReadPair, Dict[str, Any]]] = {}
    for row in rows:
        ct = str(row.get("candle_time") or "")[:12]
        if not ct:
            continue
        src, ch = normalize_source_channel(
            str(row.get("source") or ""),
            str(row.get("channel") or ""),
        )
        by_ct.setdefault(ct, {})
        pair = (src, ch)
        if pair not in by_ct[ct]:
            by_ct[ct][pair] = dict(row)

    tf = max(1, int(tf_min or 1))
    out: List[Dict[str, Any]] = []
    for ct in sorted(by_ct.keys()):
        picked: Optional[Dict[str, Any]] = None
        for pair in order:
            bar = by_ct[ct].get(pair)
            if bar is None:
                continue
            src, ch = pair
            if ticks is not None and ch == CHANNEL_WS:
                if live_cover:
                    bad = live_bar_is_garbage(
                        ticks, candle_time=ct, tf_min=tf, source=src,
                    )
                else:
                    bad = bar_is_garbage(
                        ticks, candle_time=ct, tf_min=tf, source=src,
                        missing_policy=missing_policy,
                    )
                if bad:
                    continue
            picked = dict(bar)
            break
        if picked is None:
            continue
        picked.pop("source", None)
        picked.pop("channel", None)
        out.append(picked)
    return out


def pick_bar_by_read_pairs(
    rows: Sequence[Dict[str, Any]],
    pairs: Optional[Sequence[ReadPair]] = None,
    *,
    ticks: Optional[Sequence[Dict[str, Any]]] = None,
    tf_min: int = 1,
    missing_policy: str = "hole",
    live_cover: bool = False,
) -> Optional[Dict[str, Any]]:
    """한 candle_time 묶음에서 승자 봉 1개. 없으면 None."""
    out = dedupe_by_read_pairs(
        rows, pairs, ticks=ticks, tf_min=tf_min,
        missing_policy=missing_policy, live_cover=live_cover,
    )
    return out[0] if out else None