"""
kis_trader/engine/indicator_cache.py — RSI/EMA materialized 캐시 (파라서치 가속)
================================================================================
- 종목별 IndicatorCache: RSI·EMA period 별 시리즈 1회 계산 → 조합 루프 재사용
- ws_candles DB 컬럼 materialize (RSI·EMA period 별 rsi_N / ema_N)

[ws_candles 컬럼]
- 실시간 WS 확정봉: rsi_2, rsi_3, rsi_5 (CandleAggregator 가 INSERT)
- materialize 배치: rsi_7, rsi_14, rsi_21 + ema_5,9,12,15,21,34 (env 로 period 확장)

파라서치 ④ 구간(조합×RSI/EMA 재계산) 가속용 — 실매 TRIGGER 로직 변경 없음.
"""
from __future__ import annotations

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

from kis_trader.engine.ema_trend_filter import compute_ema_series
from kis_trader.utils.env import get_env_bool, get_env_float, get_env_from_db, get_env_int


def _parse_period_list(env_key: str, fallback: str) -> List[int]:
    raw = str(get_env_from_db(env_key, fallback) or fallback).strip()
    out: List[int] = []
    for part in raw.split(","):
        part = part.strip()
        if not part:
            continue
        try:
            p = int(float(part))
            if p > 0:
                out.append(p)
        except (TypeError, ValueError):
            continue
    return sorted(set(out)) if out else [int(x) for x in fallback.split(",") if x.strip()]


def default_rsi_periods() -> List[int]:
    return _parse_period_list("PARAM_SEARCH_RSI_PERIODS", "2,3,5,7,14,21")


def default_ema_periods() -> List[int]:
    return _parse_period_list("PARAM_SEARCH_EMA_PERIODS", "5,9,12,15,21,34")


def ws_candles_materialize_rsi_periods() -> List[int]:
    """DB materialize 대상 RSI period — env ``WS_CANDLE_MATERIALIZE_RSI_PERIODS`` 우선."""
    raw = _parse_period_list(
        "WS_CANDLE_MATERIALIZE_RSI_PERIODS",
        ",".join(str(p) for p in default_rsi_periods() if p not in (2, 3, 5)),
    )
    # rsi_2/3/5 는 CandleAggregator 가 실시간 기록 — 배치는 7+ 위주
    return sorted(set(raw))


def ws_candles_materialize_ema_periods() -> List[int]:
    """DB materialize 대상 EMA period."""
    return _parse_period_list(
        "WS_CANDLE_MATERIALIZE_EMA_PERIODS",
        ",".join(str(p) for p in default_ema_periods()),
    )


def ws_candles_all_rsi_column_names() -> List[str]:
    """ws_candles SELECT/마이그레이션용 RSI 컬럼명 (rsi_2 … env period)."""
    periods = sorted(set(default_rsi_periods()) | {2, 3, 5})
    return [f"rsi_{p}" for p in periods]


def ws_candles_all_ema_column_names() -> List[str]:
    periods = ws_candles_materialize_ema_periods()
    return [f"ema_{p}" for p in periods]


def ws_candles_indicator_column_migrations() -> List[Tuple[str, str]]:
    """(컬럼명, ALTER DDL suffix) — TradeDB._migrate_add_columns 에서 사용.

    ※ TradeDB 기동 중 호출되므로 get_env_from_db(DB) 를 쓰지 않는다.
       (env 조회 → TradeDB() 재진입 → _tables_lock 데드락 → journal 로그 0줄 hang)
    """
    out: List[Tuple[str, str]] = []
    # env 기본값과 동일한 정적 period — 런타임 확장 period 는 ensure 시 information_schema 로 ADD
    migration_rsi = (7, 14, 21)
    migration_ema = (5, 9, 12, 15, 21, 34)
    for period in migration_rsi:
        out.append((
            f"rsi_{period}",
            f"DOUBLE NULL COMMENT 'RSI({period}) materialized'",
        ))
    for period in migration_ema:
        out.append((
            f"ema_{period}",
            f"DOUBLE NULL COMMENT 'EMA({period}) materialized'",
        ))
    return out


def ensure_ws_candles_indicator_columns(db) -> int:
    """ws_candles 에 누락된 rsi_N / ema_N 컬럼 ADD (1회 마이그레이션). Returns 추가된 컬럼 수."""
    added = 0
    try:
        cols = set(db.conn.get_columns("ws_candles"))
    except Exception:
        return 0
    for col, ddl in ws_candles_indicator_column_migrations():
        if col in cols:
            continue
        try:
            db.conn.execute(f"ALTER TABLE ws_candles ADD COLUMN {col} {ddl}")
            cols.add(col)
            added += 1
        except Exception:
            pass
    return added


def indicator_cache_enabled() -> bool:
    return get_env_bool("PARAM_SEARCH_INDICATOR_CACHE_ENABLED", True)


def materialize_db_on_load() -> bool:
    return get_env_bool("WS_CANDLE_MATERIALIZE_ON_LOAD", False)


class IndicatorCache:
    """단일 종목 캔들에 대한 RSI/EMA 시리즈 메모이제이션."""

    # _rsi_final / _ema_final: 해당 period 시리즈를 (DB seed + 1회 계산 병합 후)
    #   '확정'했음을 표시. 확정된 period 는 재계산 없이 캐시를 그대로 반환한다.
    #   (기존엔 rsi_at/ema_at 호출마다 compute_*_series 를 다시 돌려 O(N²) 였음 — 결과는
    #    동일했으나 파라서치 단일 백테가 봉당 전체 재계산으로 수십 초 소요되던 병목)
    __slots__ = ("closes", "_rsi", "_ema", "_rsi_final", "_ema_final")

    def __init__(self, candles: List[Dict[str, Any]]) -> None:
        self.closes: List[float] = [float(c.get("close") or 0) for c in candles]
        self._rsi: Dict[int, List[Optional[float]]] = {}
        self._ema: Dict[int, List[Optional[float]]] = {}
        self._rsi_final: set = set()
        self._ema_final: set = set()
        self._seed_from_candle_fields(candles)

    def _seed_from_candle_fields(self, candles: List[Dict[str, Any]]) -> None:
        """DB materialized 컬럼(rsi_14, ema_9 등)이 있으면 해당 period 시리즈에 반영."""
        for i, c in enumerate(candles):
            for period in default_rsi_periods():
                key = f"rsi_{period}"
                if key in c and c[key] is not None:
                    ser = self._rsi.setdefault(period, [None] * len(candles))
                    if i < len(ser):
                        try:
                            ser[i] = float(c[key])
                        except (TypeError, ValueError):
                            pass
                legacy = {2: "rsi_2", 3: "rsi_3", 5: "rsi_5"}.get(period)
                if legacy and legacy in c and c[legacy] is not None:
                    ser = self._rsi.setdefault(period, [None] * len(candles))
                    if i < len(ser) and ser[i] is None:
                        try:
                            ser[i] = float(c[legacy])
                        except (TypeError, ValueError):
                            pass
            for period in default_ema_periods():
                key = f"ema_{period}"
                if key in c and c[key] is not None:
                    ser = self._ema.setdefault(period, [None] * len(candles))
                    if i < len(ser):
                        try:
                            ser[i] = float(c[key])
                        except (TypeError, ValueError):
                            pass

    def rsi_series(self, period: int) -> List[Optional[float]]:
        p = max(1, int(period))
        # 이미 확정된 period 면 재계산 없이 캐시 반환 (핫패스 가속, 결과 불변)
        if p in self._rsi_final:
            return self._rsi[p]
        from kis_trader.engine.momentum_engine import compute_rsi_series
        computed = compute_rsi_series(self.closes, p)
        seeded = self._rsi.get(p)
        if seeded is None:
            self._rsi[p] = computed
            self._rsi_final.add(p)
            return computed
        if len(seeded) != len(computed):
            seeded = (list(seeded) + [None] * len(computed))[: len(computed)]
        # warm-up prepend 등으로 DB rsi_N 이 앞구간만 있으면 뒤 구간은 계산값으로 보완
        merged = [
            s if s is not None else b
            for s, b in zip(seeded, computed)
        ]
        self._rsi[p] = merged
        self._rsi_final.add(p)
        return merged

    def rsi_at(self, i: int, period: int) -> Optional[float]:
        ser = self.rsi_series(period)
        if 0 <= i < len(ser):
            return ser[i]
        return None

    def ema_series(self, period: int) -> List[Optional[float]]:
        p = max(1, int(period))
        # 이미 확정된 period 면 재계산 없이 캐시 반환 (핫패스 가속, 결과 불변)
        if p in self._ema_final:
            return self._ema[p]
        computed = compute_ema_series(self.closes, p)
        seeded = self._ema.get(p)
        if seeded is None:
            self._ema[p] = computed
            self._ema_final.add(p)
            return computed
        if len(seeded) != len(computed):
            seeded = (list(seeded) + [None] * len(computed))[: len(computed)]
        merged = [
            s if s is not None else b
            for s, b in zip(seeded, computed)
        ]
        self._ema[p] = merged
        self._ema_final.add(p)
        return merged

    def ema_at(self, i: int, period: int) -> Optional[float]:
        ser = self.ema_series(period)
        if 0 <= i < len(ser):
            return ser[i]
        return None


def build_indicator_cache(candles: List[Dict[str, Any]]) -> IndicatorCache:
    return IndicatorCache(candles)


def build_indicator_cache_by_code(
    codes_candles: Dict[str, List[Dict[str, Any]]],
    *,
    rsi_periods: Optional[Sequence[int]] = None,
    ema_periods: Optional[Sequence[int]] = None,
) -> Dict[str, IndicatorCache]:
    """종목별 캐시 — rsi/ema period 는 warm-up 용 (첫 접근 시 lazy 계산)."""
    if not indicator_cache_enabled():
        return {}
    out: Dict[str, IndicatorCache] = {}
    rsi_p = list(rsi_periods or default_rsi_periods())
    ema_p = list(ema_periods or default_ema_periods())
    for code, rows in (codes_candles or {}).items():
        if not rows:
            continue
        ic = IndicatorCache([dict(r) for r in rows])
        for p in rsi_p:
            ic.rsi_series(p)
        for p in ema_p:
            ic.ema_series(p)
        out[str(code)] = ic
    return out


def attach_indicator_caches_to_params(
    params: Dict[str, Any],
    codes_candles: Dict[str, List[Dict[str, Any]]],
) -> None:
    """params['_indicator_cache_by_code'] 에 종목별 캐시 주입 (in-place)."""
    if not indicator_cache_enabled():
        return
    existing = params.get("_indicator_cache_by_code")
    if isinstance(existing, dict) and existing:
        return
    params["_indicator_cache_by_code"] = build_indicator_cache_by_code(codes_candles)


def get_indicator_cache_from_params(
    params: Dict[str, Any],
    code: str,
) -> Optional[IndicatorCache]:
    by_code = params.get("_indicator_cache_by_code")
    if isinstance(by_code, dict):
        return by_code.get(str(code))
    ic = params.get("_indicator_cache")
    return ic if isinstance(ic, IndicatorCache) else None


def enrich_candles_with_materialized_fields(
    candles: List[Dict[str, Any]],
    ic: IndicatorCache,
    *,
    rsi_periods: Optional[Sequence[int]] = None,
    ema_periods: Optional[Sequence[int]] = None,
) -> None:
    """캔들 dict 에 rsi_N / ema_N 필드 부착 (DB 저장·재로드용)."""
    rsi_p = list(rsi_periods or default_rsi_periods())
    ema_p = list(ema_periods or default_ema_periods())
    for i, c in enumerate(candles):
        for p in rsi_p:
            v = ic.rsi_at(i, p)
            if v is not None:
                c[f"rsi_{p}"] = round(v, 4)
        for p in ema_p:
            v = ic.ema_at(i, p)
            if v is not None:
                c[f"ema_{p}"] = round(v, 4)


def materialize_ws_candles_batch(
    db,
    candles_by_code: Dict[str, List[Dict[str, Any]]],
    timeframe: int,
) -> int:
    """
    ws_candles 행에 rsi_N·ema_N UPDATE (배치).
    env ``WS_CANDLE_MATERIALIZE_ON_LOAD=true`` 일 때만 실행.
    Returns: 업데이트 시도 행 수.
    """
    if not materialize_db_on_load() or db is None:
        return 0
    ensure_ws_candles_indicator_columns(db)
    try:
        db_cols = set(db.conn.get_columns("ws_candles"))
    except Exception:
        return 0

    rsi_p = [p for p in ws_candles_materialize_rsi_periods() if f"rsi_{p}" in db_cols]
    ema_p = [p for p in ws_candles_materialize_ema_periods() if f"ema_{p}" in db_cols]
    if not rsi_p and not ema_p:
        return 0

    set_parts = []
    for p in rsi_p:
        set_parts.append(f"rsi_{p} = ?")
    for p in ema_p:
        set_parts.append(f"ema_{p} = ?")
    set_sql = ", ".join(set_parts)
    if not set_sql:
        return 0

    batch: List[Tuple[Any, ...]] = []
    for code, rows in (candles_by_code or {}).items():
        if not rows:
            continue
        ic = build_indicator_cache([dict(r) for r in rows])
        enrich_candles_with_materialized_fields(
            rows, ic,
            rsi_periods=rsi_p or None,
            ema_periods=ema_p or None,
        )
        for c in rows:
            ct = str(c.get("candle_time") or "")
            if not ct:
                continue
            vals: List[Any] = []
            for p in rsi_p:
                vals.append(c.get(f"rsi_{p}"))
            for p in ema_p:
                vals.append(c.get(f"ema_{p}"))
            vals.extend([str(code), int(timeframe), ct])
            batch.append(tuple(vals))

    if not batch:
        return 0

    where = "WHERE code = ? AND timeframe = ? AND candle_time = ? AND source = 'kis'"
    sql_sqlite = f"""
        UPDATE ws_candles
        SET {set_sql}, updated_at = datetime('now','localtime')
        {where}
    """
    set_mysql = ", ".join(
        [f"rsi_{p} = %s" for p in rsi_p]
        + [f"ema_{p} = %s" for p in ema_p]
    )
    sql_mysql = f"""
        UPDATE ws_candles
        SET {set_mysql}
        WHERE code = %s AND timeframe = %s AND candle_time = %s AND source = 'kis'
    """
    try:
        with db.conn:
            db.conn.executemany(sql_sqlite, batch)
    except Exception:
        try:
            db.conn.executemany(sql_mysql, batch)
        except Exception:
            return 0
    return len(batch)


def ws_candles_select_indicator_cols(db=None) -> str:
    """SELECT 절 추가 컬럼 — DB 에 실제 존재하는 rsi_N / ema_N 만."""
    names: List[str] = []
    if db is not None:
        try:
            cols = set(db.conn.get_columns("ws_candles"))
            for n in ws_candles_all_rsi_column_names() + ws_candles_all_ema_column_names():
                if n in cols:
                    names.append(n)
        except Exception:
            pass
    if not names:
        names = ws_candles_all_rsi_column_names() + ws_candles_all_ema_column_names()
    return (", " + ", ".join(names)) if names else ""


def ws_candles_has_materialized_cols(db) -> bool:
    """materialize 컬럼(ema_9 또는 ema_34 등)이 1개 이상 있으면 True."""
    try:
        cols = set(db.conn.get_columns("ws_candles"))
        for n in ws_candles_all_ema_column_names():
            if n in cols:
                return True
        for n in ("rsi_7", "rsi_14", "rsi_21"):
            if n in cols:
                return True
        return False
    except Exception:
        return False