#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
키움 ka10080 (주식분봉차트조회) 를 이용해 과거 분봉 데이터를 일괄 수집하여
``ws_candles`` 테이블에 UPSERT 한다.

목적
----
봇이 실시간으로만 캔들을 쌓으면 백테스트 기간이 짧아 파람서치 신뢰도가 낮다.
키움은 과거 봉을 수백~수천 봉까지 조회할 수 있으므로, 이 스크립트로 미리
채워두면 내일부터 바로 2주치 이상 백테스트를 돌릴 수 있다.

대상 종목
---------
기본적으로 **최근 N일 ``target_candidates_history``** 에 등장한 종목
(SCALP ∪ SHORT 합집합). ``--codes`` 로 특정 종목만 지정 가능.

대상 타임프레임
----------------
``1, 3, 15, 60`` 분 (WS_TIMEFRAMES 와 동일). 스캘핑/꼬리잡기 모두 커버.

키움 API 키
-----------
DB ``env_config`` 테이블에서 로드한다. ``KIS_MOCK`` 값에 따라 실전/모의를
자동 선택하며, 없으면 레거시 ``KIWOOM_APP_KEY`` / ``KIWOOM_APP_SECRET`` 을
사용한다.

사용 예
--------
    # 7일치, 전체 유니버스, 1/3/15/60분봉 전부 (기본값)
    python3 kis_trader/scripts/fill_kiwoom_candles.py

    # 14일치로 확장
    python3 kis_trader/scripts/fill_kiwoom_candles.py --days 14

    # 특정 종목만
    python3 kis_trader/scripts/fill_kiwoom_candles.py --codes 005930,000660

    # 1분봉/3분봉만
    python3 kis_trader/scripts/fill_kiwoom_candles.py --timeframes 1,3

    # 호출만 하고 DB 쓰지 않기 (검증용)
    python3 kis_trader/scripts/fill_kiwoom_candles.py --dry-run --codes 005930
"""

from __future__ import annotations

import argparse
import logging
import random
import sys
import time
from datetime import datetime, timedelta
from pathlib import Path
from typing import Iterable, List, Optional, Sequence

# ─── sys.path: 구봇 루트 (database.py, kis_ws.py) 접근 ─────────────────
HERE = Path(__file__).resolve()
ROOT = HERE.parents[2]            # /home/hoon/kis_bot
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from database import TradeDB                 # noqa: E402
from kis_trader.ws.kis_ws import get_kiwoom_candles_df  # noqa: E402


logger = logging.getLogger("fill_kiwoom_candles")


# ─── 상수 ─────────────────────────────────────────────────────────────
DEFAULT_TIMEFRAMES = (1, 3, 15, 60)
DEFAULT_DAYS = 7

# 장중 거래시간(분): 09:00 ~ 15:30 = 390분
# 여유분(시간외/초과 수집) 1.3배 곱해서 요청.
_MINUTES_PER_DAY = 390

# 타임프레임별 1영업일 당 기대 봉 수
def _bars_per_day(tf_min: int) -> int:
    return max(1, _MINUTES_PER_DAY // tf_min)


# ─── 키움 키 로드 ─────────────────────────────────────────────────────
def _load_kiwoom_keys(db: TradeDB) -> tuple[str, str, bool]:
    """env_config 에서 키움 앱키/시크릿/모의여부를 로드.

    우선순위
        1. KIS_MOCK=true  → KIWOOM_APP_KEY_MOCK/SECRET_MOCK
        2. KIS_MOCK=false → KIWOOM_APP_KEY_REAL/SECRET_REAL
        3. 위 둘 다 비어있으면 레거시 KIWOOM_APP_KEY/SECRET
    """
    row = db.conn.execute(
        "SELECT * FROM env_config ORDER BY id DESC LIMIT 1"
    ).fetchone()
    if not row:
        raise RuntimeError("env_config 테이블이 비어 있습니다.")
    r = dict(row)

    is_mock = str(r.get("KIS_MOCK", "")).strip().lower() in ("true", "1", "yes", "on")

    if is_mock:
        key = str(r.get("KIWOOM_APP_KEY_MOCK", "") or "").strip()
        sec = str(r.get("KIWOOM_APP_SECRET_MOCK", "") or "").strip()
        mode = "모의"
    else:
        key = str(r.get("KIWOOM_APP_KEY_REAL", "") or "").strip()
        sec = str(r.get("KIWOOM_APP_SECRET_REAL", "") or "").strip()
        mode = "실전"

    if not key or not sec:
        key = str(r.get("KIWOOM_APP_KEY", "") or "").strip()
        sec = str(r.get("KIWOOM_APP_SECRET", "") or "").strip()
        mode += "(레거시)"

    if not key or not sec:
        raise RuntimeError(
            "키움 API 키 미설정. env_config 에 KIWOOM_APP_KEY_REAL/MOCK 또는 "
            "레거시 KIWOOM_APP_KEY 가 있어야 합니다."
        )

    logger.info("🔑 키움 %s 키 로드 완료 (is_mock=%s)", mode, is_mock)
    return key, sec, is_mock


# ─── 유니버스 코드 추출 ────────────────────────────────────────────────
def _load_target_codes(
    db: TradeDB,
    *,
    days: int,
    strategies: Sequence[str],
    explicit_codes: Optional[Sequence[str]] = None,
) -> List[str]:
    """대상 종목 코드 리스트 반환.

    ``explicit_codes`` 가 있으면 그것만 사용.
    없으면 최근 ``days`` 일 ``target_candidates_history`` 에서
    ``strategies`` 에 속한 코드를 합집합으로 추출.
    """
    if explicit_codes:
        codes = [c.strip() for c in explicit_codes if c.strip()]
        logger.info("📌 --codes 지정 → %d종목", len(codes))
        return codes

    placeholders = ",".join(["%s"] * len(strategies))
    sql = f"""
        SELECT DISTINCT code
        FROM target_candidates_history
        WHERE strategy_id IN ({placeholders})
          AND event_time >= DATE_SUB(NOW(), INTERVAL %s DAY)
        ORDER BY code
    """
    params = tuple(strategies) + (int(days),)
    rows = db.conn.execute(sql, params).fetchall()
    codes = [dict(r)["code"] for r in rows]
    logger.info(
        "📌 최근 %d일 target_candidates_history ∪%s → %d종목",
        days, list(strategies), len(codes),
    )
    return codes


# ─── DB UPSERT 헬퍼 ────────────────────────────────────────────────────
# WS_CANDLE_FREEZE_ON_CONFIRM(기본 true): 존재 행 OHLCV 유지 — docs/정합성.md
_INSERT_SQL_FREEZE = """
    INSERT INTO ws_candles
        (code, timeframe, candle_time, `open`, high, low, close,
         volume, rsi_2, rsi_3, rsi_5, is_confirmed, source, updated_at)
    VALUES
        (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
    ON DUPLICATE KEY UPDATE
        candle_time=candle_time
"""

_INSERT_SQL_OVERWRITE = """
    INSERT INTO ws_candles
        (code, timeframe, candle_time, `open`, high, low, close,
         volume, rsi_2, rsi_3, rsi_5, is_confirmed, source, updated_at)
    VALUES
        (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
    ON DUPLICATE KEY UPDATE
        `open`=VALUES(`open`), high=VALUES(high), low=VALUES(low),
        close=VALUES(close), volume=VALUES(volume),
        is_confirmed=VALUES(is_confirmed), updated_at=VALUES(updated_at)
"""


def _insert_sql() -> str:
    from kis_trader.utils.env import get_env_bool

    if get_env_bool("WS_CANDLE_FREEZE_ON_CONFIRM", True):
        return _INSERT_SQL_FREEZE
    return _INSERT_SQL_OVERWRITE


def _upsert_candles(
    db: TradeDB,
    code: str,
    tf_min: int,
    df,  # pd.DataFrame
) -> int:
    """DataFrame → ws_candles UPSERT. 성공한 행 수 반환."""
    if df is None or df.empty:
        return 0

    now_str = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    rows = []
    for _, rec in df.iterrows():
        try:
            ct = str(rec["time"])[:12]
            if len(ct) < 12:
                continue
            rows.append((
                code,
                int(tf_min),
                ct,
                float(rec["open"]),
                float(rec["high"]),
                float(rec["low"]),
                float(rec["close"]),
                int(rec["volume"]),
                None, None, None,          # rsi_2, rsi_3, rsi_5 (키움 갭보정에는 없음)
                1,                          # is_confirmed (이미 확정봉)
                "kw_rest",                  # source (ws_candles.source 는 varchar(10))
                now_str,
            ))
        except Exception as e:
            logger.debug("row skip (%s %dM): %s", code, tf_min, e)

    if not rows:
        return 0

    # pymysql executemany 직접 사용
    with db.conn._lock:
        db.conn._ensure_connected()
        cur = db.conn._conn.cursor()
        cur.executemany(_insert_sql(), rows)
        db.conn._conn.commit()

    return len(rows)


# ─── 메인 루프 ────────────────────────────────────────────────────────
def run(
    *,
    days: int,
    timeframes: Sequence[int],
    strategies: Sequence[str],
    explicit_codes: Optional[Sequence[str]] = None,
    dry_run: bool = False,
    per_code_sleep: float = 0.35,
    per_tf_sleep: float = 0.25,
) -> None:
    db = TradeDB()

    kw_key, kw_sec, is_mock = _load_kiwoom_keys(db)

    codes = _load_target_codes(
        db,
        days=days,
        strategies=strategies,
        explicit_codes=explicit_codes,
    )
    if not codes:
        logger.warning("대상 종목 0개. 종료합니다.")
        return

    total_calls = len(codes) * len(timeframes)
    logger.info(
        "🚀 갭보정 시작: %d종목 × %dTF = %d호출 | 기간 %d일 | dry_run=%s",
        len(codes), len(timeframes), total_calls, days, dry_run,
    )

    cutoff_yyyymmdd = (datetime.now() - timedelta(days=days)).strftime("%Y%m%d")

    done_calls = 0
    total_rows = 0
    err_codes: list[tuple[str, int, str]] = []
    t_start = time.time()

    for ci, code in enumerate(codes, start=1):
        for tf in timeframes:
            done_calls += 1
            # 봉 수 = days × (1영업일 봉 수) × 1.3 (여유) — 페이지네이션으로 필요한만큼만
            n_req = int(_bars_per_day(tf) * days * 1.3) + 10

            try:
                df = get_kiwoom_candles_df(
                    code, tf, kw_key, kw_sec,
                    is_mock=is_mock,
                    n=n_req,
                )
            except Exception as e:
                err_codes.append((code, tf, str(e)))
                logger.warning("❌ %s %dM: %s", code, tf, e)
                time.sleep(per_tf_sleep)
                continue

            if df is None or df.empty:
                logger.debug("… %s %dM: 빈 응답", code, tf)
                time.sleep(per_tf_sleep)
                continue

            # 요청한 기간만 필터
            try:
                df2 = df[df["time"].astype(str).str[:8] >= cutoff_yyyymmdd]
            except Exception:
                df2 = df

            if dry_run:
                logger.info(
                    "🟡 [dry-run] %s %dM: 키움응답 %d봉, cutoff 이후 %d봉",
                    code, tf, len(df), len(df2),
                )
            else:
                n_saved = _upsert_candles(db, code, tf, df2)
                total_rows += n_saved
                if ci <= 3 or ci % 10 == 0 or tf == timeframes[-1]:
                    elapsed = time.time() - t_start
                    eta_s = (elapsed / done_calls) * (total_calls - done_calls)
                    logger.info(
                        "💾 [%4d/%4d] %s %2dM: +%4d봉 | 누적%6d봉 | ETA %.1f분",
                        done_calls, total_calls, code, tf, n_saved,
                        total_rows, eta_s / 60.0,
                    )

            # TF 간 짧은 sleep
            time.sleep(per_tf_sleep + random.uniform(0.0, 0.1))

        # 종목 간 sleep
        time.sleep(per_code_sleep + random.uniform(0.0, 0.15))

    elapsed_total = time.time() - t_start
    logger.info("=" * 70)
    logger.info(
        "✅ 완료: %d종목 × %dTF = %d호출 | 총 UPSERT %d봉 | 소요 %.1f분 | 실패 %d건",
        len(codes), len(timeframes), total_calls, total_rows, elapsed_total / 60.0, len(err_codes),
    )
    if err_codes:
        logger.warning("실패 목록 (최대 10건만 표시):")
        for c, tf, msg in err_codes[:10]:
            logger.warning("  %s %dM: %s", c, tf, msg[:120])


# ─── CLI ──────────────────────────────────────────────────────────────
def _parse_csv_ints(s: str) -> List[int]:
    return [int(x.strip()) for x in s.split(",") if x.strip()]


def _parse_csv_str(s: str) -> List[str]:
    return [x.strip() for x in s.split(",") if x.strip()]


def main(argv: Optional[Iterable[str]] = None) -> int:
    ap = argparse.ArgumentParser(
        description="키움 ka10080 으로 ws_candles 과거 봉 일괄 갭보정",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    ap.add_argument(
        "--days", type=int, default=DEFAULT_DAYS,
        help=f"수집 기간(일수, 기본 {DEFAULT_DAYS})",
    )
    ap.add_argument(
        "--timeframes", type=_parse_csv_ints,
        default=list(DEFAULT_TIMEFRAMES),
        help=f"타임프레임 CSV (기본 {','.join(map(str, DEFAULT_TIMEFRAMES))})",
    )
    ap.add_argument(
        "--strategies", type=_parse_csv_str,
        default=["SCALP", "SHORT"],
        help="유니버스 추출 대상 전략 (기본 SCALP,SHORT)",
    )
    ap.add_argument(
        "--codes", type=_parse_csv_str, default=None,
        help="특정 종목코드 CSV (지정 시 --strategies 무시)",
    )
    ap.add_argument(
        "--dry-run", action="store_true",
        help="키움 호출만 하고 DB 에 쓰지 않음 (검증용)",
    )
    ap.add_argument(
        "--log-level", default="INFO",
        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
        help="로그 레벨 (기본 INFO)",
    )

    args = ap.parse_args(list(argv) if argv is not None else None)

    logging.basicConfig(
        level=getattr(logging, args.log_level),
        format="%(asctime)s %(levelname)s [%(name)s] %(message)s",
        datefmt="%H:%M:%S",
    )

    run(
        days=args.days,
        timeframes=args.timeframes,
        strategies=args.strategies,
        explicit_codes=args.codes,
        dry_run=args.dry_run,
    )
    return 0


if __name__ == "__main__":
    sys.exit(main())