"""Lazy Optuna study integration used for durable trial history and pruning metadata."""

from __future__ import annotations

from pathlib import Path
from typing import Any


class OptunaStudies:
    def __init__(self, root: Path):
        self.path = (root / "optuna.sqlite3").resolve()
        self.url = f"sqlite:///{self.path}"

    def _study(self, session_id: str):
        import optuna

        optuna.logging.set_verbosity(optuna.logging.WARNING)
        return optuna.create_study(
            study_name=f"reward-tuning-{session_id}",
            storage=self.url,
            direction="maximize",
            load_if_exists=True,
            pruner=optuna.pruners.SuccessiveHalvingPruner(
                min_resource=300, reduction_factor=3, min_early_stopping_rate=0
            ),
        )

    def record(
        self,
        session_id: str,
        reward_config: dict[str, Any],
        score: float,
        eligible: bool,
        rung: int,
    ) -> int:
        """Record an externally proposed Agent config through Optuna ask/tell."""
        import optuna

        study = self._study(session_id)
        trial = study.ask()
        trial.set_user_attr("reward_config", reward_config)
        trial.set_user_attr("eligible", eligible)
        trial.set_user_attr("rung", rung)
        state = optuna.trial.TrialState.COMPLETE if eligible else optuna.trial.TrialState.PRUNED
        study.tell(trial, score if eligible else None, state=state)
        return trial.number

    def summary(self, session_id: str) -> dict[str, Any]:
        study = self._study(session_id)
        completed = [trial for trial in study.trials if trial.value is not None]
        return {
            "studyName": study.study_name,
            "trialCount": len(study.trials),
            "bestValue": max((trial.value for trial in completed), default=None),
        }