import numpy as np
from smolagents import Model

from screensuite.basebenchmark import EvaluationConfig
from screensuite.benchmark_result import BenchmarkResult
from screensuite.benchmarks.hubbasebenchmark import HubBaseBenchmark
from screensuite.benchmarks.multistep.config import AgentRunResult
from screensuite.benchmarks.multistep.gaia.config import GaiaConfig
from screensuite.benchmarks.multistep.gaia.gaia_scoring import get_question_score_gaia
from screensuite.benchmarks.multistep.generation import get_agent_responses
from screensuite.registry import registry


class GaiaBenchmark(HubBaseBenchmark[GaiaConfig, None]):
    def load(self, streaming: bool = False) -> None:
        """
        Load the dataset from Hugging Face

        Returns:
            The loaded dataset
        """
        super().load(streaming=streaming)
        assert self.dataset is not None  # type: ignore
        self.dataset = self.dataset.rename_columns({"Final answer": "reference_answer", "Question": "question"})  # type: ignore

    def evaluate(
        self,
        model: Model,
        evaluation_config: EvaluationConfig,
        env_config: None = None,
    ) -> BenchmarkResult:
        """
        Evaluate the model on the benchmark

        Args:
            model: The model to evaluate
            evaluation_config: The config for inference

        Returns:
            Evaluation results
        """
        accuracies: list[float] = []
        run_results: list[AgentRunResult] = []

        if self.dataset is None:
            self.load()

        run_results = get_agent_responses(
            self.dataset,  # type: ignore
            model,
            evaluation_config=evaluation_config,
            run_storage_folder=f"./output/gaia/{evaluation_config.run_name}",
            reformulate_responses=True,
            return_browser_url=False,
        )

        for result in run_results:
            if result is None:
                continue
            accuracies.append(get_question_score_gaia(result.answer, result.reference_answer))

        avg_accuracy = float(np.mean(accuracies)) if accuracies else 0.0

        evaluation_results = {
            "avg_accuracy": avg_accuracy,
            "proportion_missing": self._calculate_proportion_missing(run_results),
            "count_samples": len(run_results),
        }

        return BenchmarkResult(evaluation_results, "avg_accuracy")


gaia_web = GaiaBenchmark(
    name="gaia_web",
    config=GaiaConfig(),
    tags=["gaia", "multistep", "hf_dataset", "online", "to_evaluate"],
)

registry.register(gaia_web)
