from abc import ABC, abstractmethod
from typing import Any, Generator, Generic, TypeVar

from pydantic import BaseModel, Field, model_validator
from smolagents import Model
from typing_extensions import Self

from screensuite.benchmark_result import BenchmarkResult
from screensuite.benchmarks.singlestep.showdown_clicks.utils import H_SYSTEM_PROMPT

T_GroundTruthType = TypeVar("T_GroundTruthType")


class AnnotatedInput(BaseModel, Generic[T_GroundTruthType]):
    messages: list[dict[str, Any]]
    ground_truth: T_GroundTruthType


class AnnotatedOutput(BaseModel, Generic[T_GroundTruthType]):
    output: str
    ground_truth: T_GroundTruthType


class EvaluationConfig(BaseModel):
    run_name: str | None
    """The name of the run. If None, no output will be written."""
    timeout: int = Field(default=300, ge=1)
    parallel_workers: int = Field(default=1, ge=1)
    test_mode: bool = False
    max_samples_to_test: int | None = None
    use_uitars_action_space: bool = False
    use_h_action_space: bool = False
    max_image_pixels: int = 16384 * 28 * 28
    image_before_text: bool = True
    custom_system_prompt: str | None = None

    @model_validator(mode="after")
    def validate_max_samples_to_test(self) -> Self:
        if self.max_samples_to_test is not None and self.max_samples_to_test < 1:
            self.max_samples_to_test = None
        if self.max_samples_to_test is not None and self.max_samples_to_test < self.parallel_workers:
            self.parallel_workers = self.max_samples_to_test
        return self

    def model_post_init(self, context: Any) -> None:
        if self.use_h_action_space and self.use_uitars_action_space:
            raise ValueError("use_h_action_space and use_uitars_action_space cannot be True at the same time")
        if self.use_h_action_space:
            if self.custom_system_prompt is None:
                self.custom_system_prompt = H_SYSTEM_PROMPT


BaseBenchmarkConfigType = TypeVar("BaseBenchmarkConfigType", bound=BaseModel)
BaseBenchmarkEnvironmentConfigType = TypeVar("BaseBenchmarkEnvironmentConfigType")


class BaseBenchmark(Generic[BaseBenchmarkConfigType, BaseBenchmarkEnvironmentConfigType], ABC):
    def __init__(self, name: str, config: BaseBenchmarkConfigType, tags: list[str]):
        self.name = name
        self.config = config
        self.tags = tags

    @abstractmethod
    def load(self) -> Any: ...

    @abstractmethod
    def evaluate(
        self,
        model: Model,
        evaluation_config: EvaluationConfig,
        env_config: BaseBenchmarkEnvironmentConfigType = ...,
    ) -> BenchmarkResult: ...

    def _get_annotated_input_from_sample(
        self, sample: Any, evaluation_config: EvaluationConfig
    ) -> AnnotatedInput | Generator[AnnotatedInput, None, None]:
        """
        This method is used to transform the sample into an AnnotatedInput object containing a list of messages and a ground truth.
        """
        raise NotImplementedError("Transforming samples is not implemented for this benchmark")
