from pydantic import BaseModel
from smolagents import Model

from screensuite.basebenchmark import BaseBenchmark, EvaluationConfig
from screensuite.benchmark_result import BenchmarkResult
from screensuite.registry import registry


class ExampleBenchmarkConfig(BaseModel):
    """
    A configuration for the example benchmark.
    """

    pass


class ExampleBenchmark(BaseBenchmark[ExampleBenchmarkConfig, None]):
    """
    An example benchmark that demonstrates how to implement a benchmark.
    """

    def __init__(self):
        super().__init__(
            name="example_benchmark",
            config=ExampleBenchmarkConfig(),
            tags=["example", "demo"],
        )

    def load(self) -> None:
        """
        Load the benchmark data.
        """
        # In a real benchmark, this would load data from a file or API
        self.data = {
            "examples": [
                {"input": "What is 2+2?", "expected": "4"},
                {"input": "What is the capital of France?", "expected": "Paris"},
            ]
        }

    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 inference configuration
        Returns:
            A dictionary containing the evaluation results
        """
        if not hasattr(self, "data"):
            self.load()

        correct = 0
        total = len(self.data["examples"])

        for example in self.data["examples"]:
            response = model(messages=[{"role": "user", "content": example["input"]}])
            if response.content is not None and response.content.strip().lower() == example["expected"].lower():
                correct += 1

        accuracy = correct / total if total > 0 else 0

        results = {"accuracy": accuracy, "correct": correct, "total": total}

        return BenchmarkResult(results, "accuracy")


# Register the example benchmark
example_benchmark = ExampleBenchmark()
registry.register(example_benchmark)
