# Optuna-powered attack hyperparameter tuning.


from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Optional

logger = logging.getLogger(__name__)


class OptunaTuner:

    def __init__(self, n_trials: int = 50, timeout_seconds: Optional[int] = None) -> None:
        self.n_trials = n_trials
        self.timeout_seconds = timeout_seconds

    def tune(self, model, attack_classes: list, verbose: bool = True) -> dict:
        try:
            import optuna
            optuna.logging.set_verbosity(
                optuna.logging.INFO if verbose else optuna.logging.WARNING
            )
        except ImportError:
            logger.warning("optuna not installed — returning default params")
            return {}

        best_params = {}

        for attack_cls in attack_classes:
            attack_name = getattr(attack_cls, "attack_name", attack_cls.__name__)
            param_space = self._get_param_space(attack_cls)
            if not param_space:
                continue

            logger.info("Tuning %s over %d trials...", attack_name, self.n_trials)

            def objective(trial, _attack_cls=attack_cls, _model=model, _space=param_space):
                params = self._sample_params(trial, _space)
                attack = _attack_cls(**{k: v for k, v in params.items() if k in self._get_init_params(_attack_cls)})
                result = attack.run(_model, **{k: v for k, v in params.items() if k not in self._get_init_params(_attack_cls)})
                return self._score_result(result)

            study = optuna.create_study(direction="maximize")
            study.optimize(objective, n_trials=self.n_trials, timeout=self.timeout_seconds)
            best_params[attack_name] = study.best_params
            logger.info("Best params for %s: %s (score=%.3f)", attack_name, study.best_params, study.best_value)

        return best_params

    def _get_param_space(self, attack_cls) -> dict:
        attack_name = getattr(attack_cls, "attack_name", "")

        spaces = {
            "simba_token": {
                "n_iters": ("int", 10, 100),
            },
            "nes_gradient_estimate": {
                "n_samples": ("int", 5, 30),
                "n_iters": ("int", 10, 50),
                "sigma": ("float", 0.01, 0.5),
            },
            "hopskipjump": {
                "n_iters": ("int", 10, 50),
            },
            "model_extraction": {
                "depth": ("categorical", ["quick", "standard"]),
            },
        }
        return spaces.get(attack_name, {})

    def _sample_params(self, trial, space: dict) -> dict:
        params = {}
        for name, spec in space.items():
            if spec[0] == "int":
                params[name] = trial.suggest_int(name, spec[1], spec[2])
            elif spec[0] == "float":
                params[name] = trial.suggest_float(name, spec[1], spec[2])
            elif spec[0] == "categorical":
                params[name] = trial.suggest_categorical(name, spec[1])
        return params

    def _get_init_params(self, attack_cls) -> list[str]:
        import inspect
        sig = inspect.signature(attack_cls.__init__)
        return [p for p in sig.parameters if p != "self"]

    def _score_result(self, result) -> float:
        from sichgate.schema import AttackResult
        if not isinstance(result, AttackResult):
            return 0.0
        severity_scores = {"critical": 1.0, "high": 0.75, "medium": 0.5, "low": 0.1}
        status_scores = {"fail": 1.0, "warning": 0.5, "pass": 0.0}
        s = severity_scores.get(result.severity, 0.0) * status_scores.get(result.status, 0.0)
        # Bonus if telemetry shows completed
        if result.execution_telemetry.status == "completed":
            s += 0.1
        return min(s, 1.0)


def run_attacks(model_path, attack_classes, adapter_type, depth, domain_data):
    with console.status("Running attacks..."):
        summary = runner.run(
            model_path=model_path,
            attack_classes=attack_classes,
            adapter_type=adapter_type,
            depth=depth,
            domain_data_path=domain_data,
        )
    return summary
