#!/usr/bin/env python3
"""kis_trader/backtest/optuna_dart.py — DART Optuna TPE."""
from __future__ import annotations

import logging
from datetime import datetime
from typing import Any, Dict, Optional

import optuna
from optuna.samplers import TPESampler

from database import TradeDB
from kis_trader.backtest.optuna_common import (
    announce_optuna_json_path,
    resolve_optuna_storage_url,
)
from kis_trader.backtest.param_search_dart import (
    _dart_grids,
    apply_params_to_db,
    evaluate_dart_param_combo,
)
from kis_trader.engine import dart_engine as de

logger = logging.getLogger("param_search_optuna")
_FAIL = -1e18


def prepare_dart_search_context(start: str, end: str, mode: str) -> Optional[Dict[str, Any]]:
    grids = _dart_grids()
    if mode not in grids:
        logger.error("DART mode: %s (fast/coarse/fine)", mode)
        return None
    db = TradeDB()
    try:
        env_row = db.get_merged_env_snapshot() or {}
    finally:
        db.close()
    return {
        "start": start,
        "end": end,
        "mode": mode,
        "grid": grids[mode],
        "env_row": env_row,
        "base": de.get_dart_defaults_from_db(env_row=env_row),
    }


def run_dart_optuna(
    ctx: Dict[str, Any],
    *,
    n_trials: int,
    storage_url: Optional[str] = None,
    study_name: Optional[str] = None,
    min_trades: int = 1,
    sampler_name: str = "tpe",
    seed: int = 42,
    show_progress: bool = False,
    **_kwargs: Any,
) -> optuna.Study:
    grid: Dict[str, Any] = ctx["grid"]
    storage = storage_url or resolve_optuna_storage_url()
    name = study_name or "dart_%s_%s" % (
        ctx["mode"], datetime.now().strftime("%Y%m%d_%H%M%S"),
    )
    sampler = TPESampler(seed=seed) if sampler_name != "random" else optuna.samplers.RandomSampler(seed=seed)
    study = optuna.create_study(
        study_name=name,
        storage=storage,
        direction="maximize",
        sampler=sampler,
        load_if_exists=True,
    )

    def objective(trial: optuna.Trial) -> float:
        params = dict(ctx["base"])
        for axis, values in grid.items():
            params[axis] = trial.suggest_categorical(axis, list(values))
        res = evaluate_dart_param_combo(
            params, start=ctx["start"], end=ctx["end"], env_row=ctx["env_row"],
        )
        n = int(res.get("trade_count") or 0)
        pnl = float(res.get("total_pnl") or 0)
        trial.set_user_attr("trade_count", n)
        trial.set_user_attr("total_pnl", pnl)
        trial.set_user_attr("win_rate", float(res.get("win_rate") or 0))
        if n < min_trades:
            return _FAIL
        return pnl

    study.optimize(objective, n_trials=n_trials, show_progress_bar=show_progress)
    announce_optuna_json_path("dart", study)
    return study


def apply_best_dart_trial(study: optuna.Study) -> bool:
    try:
        best = study.best_trial
    except Exception:
        return False
    if float(best.value or _FAIL) <= 0:
        logger.warning("DART best PnL<=0 — DB 미적용")
        return False
    apply_params_to_db(dict(best.params))
    return True