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

UI/CLI ``CANDLE_SOURCE``:
  - 빈값(「기본」): LIVE_TICK_PROVIDER 순서로 candle_time 디듑
    kiwoom 메인 → kiwoom, kis, rest, rollup_1m
    kis 메인    → kis, kiwoom, rest, rollup_1m
  - ``kis`` / ``kiwoom``: 해당 소스만 (디버그·비교용)
"""
from __future__ import annotations

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

from kis_trader.utils.env import get_env_from_db

# 실매 CandleAggregator._ALL_CANDLE_SOURCES 와 동일
BT_WS_CANDLE_SOURCES: Tuple[str, ...] = ("kiwoom", "kis", "rest", "rollup_1m")


def resolve_bt_candle_source_override() -> str:
    """'' = 병합 모드, 'kis'|'kiwoom' = 단일 소스."""
    raw = os.environ.get("CANDLE_SOURCE")
    if raw is None:
        try:
            raw = get_env_from_db("CANDLE_SOURCE", "")
        except Exception:
            raw = ""
    s = (str(raw or "")).strip().lower()
    if s in ("kis", "kiwoom"):
        return s
    return ""


def live_candle_source_order() -> Tuple[str, ...]:
    """실매 get_candles 병합 순서 (LIVE_TICK_PROVIDER)."""
    try:
        provider = (
            get_env_from_db("LIVE_TICK_PROVIDER", "kiwoom") or "kiwoom"
        ).strip().lower()
    except Exception:
        provider = "kiwoom"
    if provider == "kis":
        return ("kis", "kiwoom", "rest", "rollup_1m")
    return ("kiwoom", "kis", "rest", "rollup_1m")


def resolve_bt_candle_source_order() -> Tuple[str, ...]:
    """백테/Optuna 조회에 쓸 소스 순서."""
    override = resolve_bt_candle_source_override()
    if override:
        return (override,)
    return live_candle_source_order()


def resolve_bt_candle_source_label() -> str:
    """상태 표시용 — 'KIS'|'KIWOOM'|'LIVE(kiwoom)' 등."""
    override = resolve_bt_candle_source_override()
    if override:
        return override.upper()
    try:
        provider = (
            get_env_from_db("LIVE_TICK_PROVIDER", "kiwoom") or "kiwoom"
        ).strip().lower()
    except Exception:
        provider = "kiwoom"
    return f"LIVE({provider})"


def dedupe_candle_rows(
    rows: Sequence[Dict[str, Any]],
    source_order: Optional[Sequence[str]] = None,
) -> List[Dict[str, Any]]:
    """candle_time 기준 디듑 — 앞 소스 우선. ``source`` 컬럼은 결과에서 제거."""
    order = tuple(source_order or resolve_bt_candle_source_order())
    rank = {s: i for i, s in enumerate(order)}
    seen: Dict[str, Tuple[int, Dict[str, Any]]] = {}
    for row in rows:
        ct = str(row.get("candle_time") or "")[:12]
        if not ct:
            continue
        src = str(row.get("source") or "kis").strip().lower()
        pri = rank.get(src, 999)
        prev = seen.get(ct)
        if prev is None or pri < prev[0]:
            seen[ct] = (pri, dict(row))
    out: List[Dict[str, Any]] = []
    for ct in sorted(seen.keys()):
        d = seen[ct][1]
        d.pop("source", None)
        out.append(d)
    return out


def _source_in_sql(sources: Sequence[str]) -> Tuple[str, List[str]]:
    if len(sources) == 1:
        return " AND source=%s", [sources[0]]
    ph = ",".join(["%s"] * len(sources))
    return f" AND source IN ({ph})", list(sources)


def list_ws_candle_codes(
    db,
    timeframe: int,
    start_key: str,
    end_key: str,
    *,
    market: Optional[str] = None,
) -> List[str]:
    """기간 내 종목 코드 — 선택된 소스 기준 DISTINCT."""
    sources = resolve_bt_candle_source_order()
    src_sql, src_params = _source_in_sql(sources)
    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]]:
    """단일 종목·기간 봉 로드 — 소스 필터/병합 적용."""
    sources = resolve_bt_candle_source_order()
    confirmed_sql = " AND is_confirmed=1" if confirmed_only else ""
    mk = (market or "").strip().upper()

    if len(sources) == 1:
        cols = f"candle_time, open, high, low, close, volume{peak_sel}{extra_select}"
        src = sources[0]
        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
                + " AND source=%s ORDER BY candle_time ASC",
                [int(timeframe), code, mk, start_key, end_key, src],
            ).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
                + " AND source=%s ORDER BY candle_time ASC",
                [int(timeframe), code, start_key, end_key, src],
            ).fetchall()
        return [dict(r) for r in rows]

    cols = f"candle_time, open, high, low, close, volume, source{peak_sel}{extra_select}"
    src_sql, src_params = _source_in_sql(sources)
    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()
    return dedupe_candle_rows([dict(r) for r in rows], sources)


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 []
    sources = resolve_bt_candle_source_order()
    confirmed_sql = " AND is_confirmed=1" if confirmed_only else ""
    fetch_limit = max(int(limit) * max(len(sources), 1), int(limit) + 50)

    if len(sources) == 1:
        cols = f"candle_time, open, high, low, close, volume{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
            + " AND source=%s ORDER BY candle_time DESC LIMIT %s",
            [int(timeframe), code, before_candle_time, sources[0], fetch_limit],
        ).fetchall()
        bars = [dict(r) for r in reversed(rows)]
        return bars[-limit:] if len(bars) > limit else bars

    cols = f"candle_time, open, high, low, close, volume, source{peak_sel}{extra_select}"
    src_sql, src_params = _source_in_sql(sources)
    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()
    bars = dedupe_candle_rows([dict(r) for r in rows], sources)
    return bars[-limit:] if len(bars) > limit else bars