"""ensemble_weights 테이블 IO. 기본 가중치 (chronos 0.6, lgbm 0.4)."""
from __future__ import annotations

from dataclasses import dataclass

from sqlalchemy import text

from app.db.connection import get_engine


@dataclass
class EnsembleWeights:
    code: str
    horizon: int
    w_chronos: float
    w_lgbm: float


DEFAULT_W_CHRONOS = 0.6
DEFAULT_W_LGBM = 0.4


def load_weights(code: str, horizon: int) -> EnsembleWeights:
    eng = get_engine()
    with eng.connect() as conn:
        row = conn.execute(
            text(
                "SELECT w_chronos, w_lgbm FROM ensemble_weights "
                "WHERE code = :code AND horizon = :h"
            ),
            {"code": code, "h": horizon},
        ).first()
    if not row:
        return EnsembleWeights(code, horizon, DEFAULT_W_CHRONOS, DEFAULT_W_LGBM)
    return EnsembleWeights(code, horizon, float(row[0]), float(row[1]))


def upsert_weights(
    code: str,
    horizon: int,
    w_chronos: float,
    w_lgbm: float,
    *,
    hit_rate_chronos: float | None = None,
    hit_rate_lgbm: float | None = None,
    sample_count: int | None = None,
) -> None:
    eng = get_engine()
    with eng.begin() as conn:
        conn.execute(
            text(
                """
                INSERT INTO ensemble_weights
                  (code, horizon, w_chronos, w_lgbm, hit_rate_chronos, hit_rate_lgbm, sample_count, updated_at)
                VALUES
                  (:code, :h, :wc, :wl, :hc, :hl, :n, NOW())
                ON CONFLICT (code, horizon) DO UPDATE SET
                  w_chronos = EXCLUDED.w_chronos,
                  w_lgbm    = EXCLUDED.w_lgbm,
                  hit_rate_chronos = EXCLUDED.hit_rate_chronos,
                  hit_rate_lgbm    = EXCLUDED.hit_rate_lgbm,
                  sample_count     = EXCLUDED.sample_count,
                  updated_at = NOW()
                """
            ),
            {
                "code": code,
                "h": horizon,
                "wc": float(w_chronos),
                "wl": float(w_lgbm),
                "hc": hit_rate_chronos,
                "hl": hit_rate_lgbm,
                "n": sample_count,
            },
        )