#!/usr/bin/env python3
"""
kis_trader/backtest/optuna_mode_refine_runner.py — 1·2차 TPE 연쇄 (기간=폼 start/end)
=====================================================================================
1차: 넓은 Grid TPE → JSON
2차: 1차 JSON 밴드로 Grid 축소 + 동일 기간 TPE

웹: start_optuna_job(단일) · run_optuna_4strat_tpe_seq.sh(순차) → subprocess 본 스크립트.
"""
from __future__ import annotations

import argparse
import json
import logging
import os
import subprocess
import sys
import time
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional

ROOT = Path(__file__).resolve().parents[2]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

logger = logging.getLogger("optuna_mode_refine")


def _py() -> str:
    cand = ROOT / ".venv" / "bin" / "python3"
    if cand.is_file():
        return str(cand)
    return sys.executable


def _read_json(path: str) -> Dict[str, Any]:
    return json.loads(Path(path).read_text(encoding="utf-8"))


def _load_phase1_data(
    *,
    strategy: str,
    phase1_json: Optional[str] = None,
    phase1_study: Optional[str] = None,
) -> tuple:
    """1차 narrow 입력. DB(study) 우선, 없으면 JSON 파일. (data, source_label)."""
    study = str(phase1_study or "").strip()
    if study:
        from kis_trader.backtest.optuna_study_store import load_payload_dict, payload_has_rows

        data = load_payload_dict(study)
        if data and payload_has_rows(data):
            logger.info("📌 1차 payload DB study=%s", study)
            return data, f"db:{study}"
        logger.warning("⚠️ DB payload 없음 study=%s — JSON 폴백 시도", study)
    path = str(phase1_json or "").strip()
    if path and Path(path).is_file():
        return _read_json(path), path
    if study:
        found = _find_json_by_study(study, strategy)
        if found and Path(found).is_file():
            return _read_json(found), found
    raise FileNotFoundError(
        f"1차 payload 없음 (study={study or '-'}, json={path or '-'})"
    )


def _persist_phase1_to_db(data: Dict[str, Any], *, job_id: str, study_name: str) -> None:
    """1차 JSON → optuna_study_result (다른 PC 2차용)."""
    try:
        from kis_trader.backtest.optuna_study_store import ingest_out_data

        out = ingest_out_data(data, job_id=job_id, replace=False)
        logger.info("📌 1차 DB 저장 study=%s ok=%s", study_name, out.get("ok"))
    except Exception as exc:
        logger.warning("⚠️ 1차 DB 저장 실패 study=%s: %s", study_name, exc)


def _write_state(path: Path, state: Dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(state, indent=2, ensure_ascii=False), encoding="utf-8")


def _find_json_by_study(study_name: str, strategy: str) -> Optional[str]:
    from kis_trader.backtest.tail_param_search import _results_dir_for_write

    out_dir = Path(_results_dir_for_write())
    if not out_dir.is_dir():
        return None
    cands: List[Path] = []
    patterns = [
        f"optuna_{strategy}_*.json",
        f"optuna_{strategy[:4]}_*.json",
        "optuna_*.json",
    ]
    seen: set = set()
    for pat in patterns:
        for p in out_dir.glob(pat):
            if str(p) in seen:
                continue
            seen.add(str(p))
            try:
                d = json.loads(p.read_text(encoding="utf-8"))
            except Exception:
                continue
            if str(d.get("optuna_study_name") or "") == study_name:
                cands.append(p)
    if not cands:
        return None
    cands.sort(key=lambda p: p.stat().st_mtime, reverse=True)
    return str(cands[0])


def _build_optuna_cmd(
    *,
    strategy: str,
    mode: str,
    start: str,
    end: str,
    trials: int,
    study_name: str,
    sort_by: str,
    min_trades: int,
    hist_src: str,
    entry_mode: Optional[str] = None,
    sl_mode: Optional[str] = None,
    ob_mode: Optional[str] = None,
    symbol: Optional[str] = None,
    candle_source: Optional[str] = None,
    tick_source: Optional[str] = None,
    ob_source: Optional[str] = None,
    study_trials: Optional[int] = None,
) -> List[str]:
    cmd = [
        _py(), "-u",
        str(ROOT / "kis_trader" / "backtest" / "param_search_optuna.py"),
        "--strategy", strategy,
        "--mode", mode,
        "--start", start,
        "--end", end,
        "--trials", str(int(trials)),
        "--min_trades", str(int(min_trades)),
        "--min_win_rate", "0",
        "--min_pf", "0",
        "--no-progress",
        "--study-name", study_name,
        "--sort-by", sort_by,
        "--universe-history-source", hist_src,
    ]
    if strategy == "tail" and entry_mode:
        cmd.extend(["--entry-mode", entry_mode])
    if strategy == "breakout":
        cmd.extend(["--sl-mode", sl_mode or "fixed"])
        cmd.extend(["--orderbook-filter", ob_mode or "off"])
    if symbol and strategy == "us_momentum":
        cmd.extend(["--symbol", symbol])
    if candle_source:
        cmd.extend(["--candle-source", candle_source])
    if tick_source:
        cmd.extend(["--tick-source", tick_source])
    if ob_source:
        cmd.extend(["--ob-source", ob_source])
    if study_trials and int(study_trials) > 0:
        cmd.extend(["--study-trials", str(int(study_trials))])
    return cmd


def _run_phase(
    cmd: List[str],
    env: Dict[str, str],
    log_path: Path,
) -> int:
    log_path.parent.mkdir(parents=True, exist_ok=True)
    with open(log_path, "w", encoding="utf-8") as log_f:
        log_f.write("CMD: %s\n\n" % " ".join(cmd))
        log_f.flush()
        proc = subprocess.Popen(
            cmd,
            cwd=str(ROOT),
            env=env,
            stdout=log_f,
            stderr=subprocess.STDOUT,
        )
        return int(proc.wait())


def run_mode_refine(
    *,
    job_id: str,
    strategy: str,
    mode: str,
    start: str,
    end: str,
    trials: int,
    sort_by: str,
    min_trades: int,
    hist_src: str,
    entry_mode: Optional[str] = None,
    sl_mode: Optional[str] = None,
    ob_mode: Optional[str] = None,
    symbol: Optional[str] = None,
    candle_source: Optional[str] = None,
    tick_source: Optional[str] = None,
    ob_source: Optional[str] = None,
    study_trials: Optional[int] = None,
    skip_phase1: bool = False,
    phase1_json: Optional[str] = None,
    phase1_study: Optional[str] = None,
) -> Dict[str, Any]:
    from kis_trader.backtest.optuna_grid_narrow import (
        build_and_write_narrow_grid,
        resolve_refine_phase2_trials,
    )
    from kis_trader.backtest.optuna_mode_combo import resolve_mode_pool_kind

    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    state_path = ROOT / "logs" / f"{job_id}_refine_state.json"
    narrow_path = ROOT / "logs" / f"{job_id}_narrow_grid.json"
    state: Dict[str, Any] = {
        "job_id": job_id,
        "strategy": strategy,
        "mode": mode,
        "start": start,
        "end": end,
        "phase": "phase1",
        "pool_kind": resolve_mode_pool_kind(),
    }
    _write_state(state_path, state)

    env = os.environ.copy()
    env["PYTHONUNBUFFERED"] = "1"
    env["PYTHONPATH"] = str(ROOT) + (os.pathsep + env.get("PYTHONPATH", ""))
    env["BACKTEST_UNIVERSE_HISTORY_SOURCE"] = hist_src
    env["PARAM_SEARCH_OPTUNA_REPORT_MIN_TRADES"] = str(min_trades)
    env["SORT_BY"] = sort_by
    env["OPTUNA_WEB_JOB_ID"] = str(job_id)
    env.pop("OPTUNA_GRID_NARROW_JSON", None)

    p1_study_arg = str(phase1_study or "").strip()
    
    extra = ""
    if strategy == "tail" and entry_mode:
        extra = f"_{entry_mode}"
    elif strategy == "breakout":
        from kis_trader.backtest.optuna_breakout_tpe_space import breakout_tpe_study_extra
        extra = f"_{breakout_tpe_study_extra(sl_mode or 'fixed', ob_mode or 'off')}"

    p1_study = f"{strategy}{extra}_{mode}_refine1_{start.replace('-', '')}_{end.replace('-', '')}_{ts}"
    p1_log = ROOT / "logs" / f"optuna_refine1_{job_id}.log"
    p1_json = str(phase1_json or "").strip()

    if skip_phase1 and (p1_study_arg or (p1_json and Path(p1_json).is_file())):
        logger.info("⏭ 1차 스킵 — study=%s json=%s", p1_study_arg or "-", p1_json or "-")
        data, src = _load_phase1_data(
            strategy=strategy,
            phase1_json=p1_json or None,
            phase1_study=p1_study_arg or None,
        )
        p1_study = str(
            p1_study_arg
            or data.get("optuna_study_name")
            or data.get("study_name")
            or p1_study
        ).strip()
        state["phase1_study"] = p1_study
        state["phase1_json"] = p1_json or src
        state["phase1_source"] = src
        state["phase1_skipped"] = True
    else:
        cmd1 = _build_optuna_cmd(
            strategy=strategy,
            mode=mode,
            start=start,
            end=end,
            trials=trials,
            study_name=p1_study,
            sort_by=sort_by,
            min_trades=min_trades,
            hist_src=hist_src,
            entry_mode=entry_mode,
            sl_mode=sl_mode,
            ob_mode=ob_mode,
            symbol=symbol,
            candle_source=candle_source,
            tick_source=tick_source,
            ob_source=ob_source,
            study_trials=study_trials,
        )
        state["phase1_study"] = p1_study
        state["phase1_log"] = str(p1_log)
        state["phase"] = "phase1"
        _write_state(state_path, state)
        logger.info("▶ 1차 TPE 시작 study=%s %s~%s trials=%s", p1_study, start, end, trials)
        rc1 = _run_phase(cmd1, env, p1_log)
        state["phase1_exit"] = rc1
        state["phase1_log"] = str(p1_log)
        state["phase1_study"] = p1_study
        _write_state(state_path, state)
        if rc1 != 0:
            state["phase"] = "error"
            state["error"] = f"phase1 exit {rc1}"
            _write_state(state_path, state)
            return state
        p1_json = _find_json_by_study(p1_study, strategy) or ""
        if not p1_json or not Path(p1_json).is_file():
            state["phase"] = "error"
            state["error"] = "phase1 JSON not found"
            _write_state(state_path, state)
            return state
        state["phase1_json"] = p1_json
        state["phase1_study"] = p1_study
        print(f"OPTUNA_PHASE1_STUDY={p1_study}", flush=True)
        try:
            _persist_phase1_to_db(_read_json(p1_json), job_id=job_id, study_name=p1_study)
        except Exception:
            pass

    try:
        data = _load_phase1_data(
            strategy=strategy,
            phase1_json=p1_json or None,
            phase1_study=str(state.get("phase1_study") or p1_study_arg or "").strip() or None,
        )[0]
    except FileNotFoundError as exc:
        state["phase"] = "error"
        state["error"] = str(exc)
        _write_state(state_path, state)
        return state

    if state.get("phase1_skipped"):
        p1_study = str(state.get("phase1_study") or p1_study).strip()
        state["phase1_study"] = p1_study
    data["mode_pool_kind"] = resolve_mode_pool_kind()
    narrow_meta = build_and_write_narrow_grid(
        data, str(narrow_path), mode=mode, strategy=strategy,
    )
    state["narrow_grid"] = narrow_meta
    state["narrow_path"] = str(narrow_path)

    p2_trials = resolve_refine_phase2_trials(trials)
    p2_study = f"{strategy}{extra}_{mode}_refine2_{start.replace('-', '')}_{end.replace('-', '')}_{ts}"
    p2_log = ROOT / "logs" / f"optuna_refine2_{job_id}.log"
    env2 = dict(env)
    env2["OPTUNA_GRID_NARROW_JSON"] = str(narrow_path)

    cmd2 = _build_optuna_cmd(
        strategy=strategy,
        mode=mode,
        start=start,
        end=end,
        trials=p2_trials,
        study_name=p2_study,
        sort_by=sort_by,
        min_trades=min_trades,
        hist_src=hist_src,
        entry_mode=entry_mode,
        sl_mode=sl_mode,
        ob_mode=ob_mode,
        symbol=symbol,
        candle_source=candle_source,
        tick_source=tick_source,
        ob_source=ob_source,
        study_trials=study_trials,
    )
    logger.info(
        "▶ 2차 TPE(밴드축소) study=%s trials=%s narrow=%s",
        p2_study, p2_trials, narrow_path,
    )
    state["phase"] = "phase2"
    state["phase2_study"] = p2_study
    state["phase2_log"] = str(p2_log)
    state["phase2_trials"] = p2_trials
    _write_state(state_path, state)
    rc2 = _run_phase(cmd2, env2, p2_log)
    state["phase2_exit"] = rc2
    state["phase2_log"] = str(p2_log)
    state["phase2_study"] = p2_study
    state["phase2_trials"] = p2_trials
    p2_json = _find_json_by_study(p2_study, strategy) or ""
    if p2_json:
        state["phase2_json"] = p2_json
        state["result_json"] = p2_json
    state["phase"] = "done" if rc2 == 0 else "error"
    if rc2 != 0:
        state["error"] = f"phase2 exit {rc2}"
    state["finished_at"] = datetime.now().isoformat(timespec="seconds")
    _write_state(state_path, state)
    return state


def main() -> int:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
    )
    ap = argparse.ArgumentParser(description="Optuna 1·2차 TPE 연쇄 (밴드 축소 2차)")
    ap.add_argument("--job-id", required=True)
    ap.add_argument("--strategy", required=True)
    ap.add_argument("--mode", default="tpe")
    ap.add_argument("--start", required=True)
    ap.add_argument("--end", required=True)
    ap.add_argument("--trials", type=int, default=200)
    ap.add_argument("--sort-by", default="score")
    ap.add_argument("--min-trades", type=int, default=1)
    ap.add_argument("--universe-history-source", default="kiwoom")
    ap.add_argument("--entry-mode", default="")
    ap.add_argument("--sl-mode", default="fixed")
    ap.add_argument("--ob-mode", default="off")
    ap.add_argument("--symbol", default="")
    ap.add_argument("--candle-source", default="")
    ap.add_argument("--tick-source", default="")
    ap.add_argument("--ob-source", default="")
    ap.add_argument("--study-trials", type=int, default=0)
    ap.add_argument("--skip-phase1", action="store_true")
    ap.add_argument("--phase1-json", default="")
    ap.add_argument("--phase1-study", default="", help="1차 study — MariaDB payload (다른 PC 2차)")
    args = ap.parse_args()

    st = run_mode_refine(
        job_id=str(args.job_id),
        strategy=str(args.strategy).strip().lower(),
        mode=str(args.mode).strip().lower() or "tpe",
        start=str(args.start).strip(),
        end=str(args.end).strip(),
        trials=max(1, int(args.trials)),
        sort_by=str(args.sort_by or "score"),
        min_trades=max(1, int(args.min_trades)),
        hist_src=str(args.universe_history_source or "kiwoom").strip().lower(),
        entry_mode=str(args.entry_mode or "").strip() or None,
        sl_mode=str(args.sl_mode or "fixed").strip() or "fixed",
        ob_mode=str(args.ob_mode or "off").strip() or "off",
        symbol=str(args.symbol or "").strip().upper() or None,
        candle_source=str(args.candle_source or "").strip().lower() or None,
        tick_source=str(args.tick_source or "").strip().lower() or None,
        ob_source=str(args.ob_source or "").strip().lower() or None,
        study_trials=int(args.study_trials) if int(args.study_trials or 0) > 0 else None,
        skip_phase1=bool(args.skip_phase1),
        phase1_json=str(args.phase1_json or "").strip() or None,
        phase1_study=str(args.phase1_study or "").strip() or None,
    )
    if st.get("phase") == "done":
        logger.info("✅ 1·2차 완료 result=%s", st.get("result_json"))
        return 0
    logger.error("❌ 1·2차 실패: %s", st.get("error"))
    return 1


if __name__ == "__main__":
    raise SystemExit(main())