"""Base classes for benchmark definitions.

This module defines the abstract interfaces for benchmarks in the EB1
evaluation system. All benchmark loaders must implement these interfaces.
"""

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Iterator


class BenchmarkDomain(Enum):
    """High-level domain categories for benchmarks."""

    CODING = "coding"
    MATH = "math"
    KNOWLEDGE = "knowledge"
    INSTRUCTION = "instruction"
    REASONING = "reasoning"
    COMMONSENSE = "commonsense"


class MetricType(Enum):
    """Primary metric type for benchmark evaluation."""

    EXACT_MATCH = "exact_match"  # String equality
    PASS_AT_K = "pass@k"  # Code execution
    F1 = "f1"  # Token-level F1
    LLM_JUDGE = "llm_judge"  # LLM-as-judge scoring
    CONSTRAINT = "constraint"  # Constraint satisfaction (IFEval)
    NUMERIC = "numeric"  # Numeric answer matching


class DifficultyLevel(Enum):
    """Difficulty levels for slicing analysis."""

    TRIVIAL = "trivial"  # All models correct
    EASY = "easy"  # Most models correct
    MEDIUM = "medium"  # Mixed performance
    HARD = "hard"  # Few models correct
    VERY_HARD = "very_hard"  # Very few correct
    IMPOSSIBLE = "impossible"  # No models correct
    UNKNOWN = "unknown"  # Difficulty not computed


@dataclass
class BenchmarkExample:
    """A single evaluation example from a benchmark.

    Attributes:
        id: Unique identifier for this example.
        query: The input prompt/question to send to the model.
        reference: Ground truth answer(s). Can be a single string or list
            for multiple valid answers.
        metadata: Arbitrary metadata (task, difficulty, domain, source, etc.)
        system_prompt: Optional system prompt for the model.
        context: Optional context for RAG-style benchmarks.
        constraints: Optional list of constraints (for IFEval).
        test_cases: Optional test cases for code execution benchmarks.
    """

    id: str
    query: str
    reference: str | list[str]
    metadata: dict[str, Any] = field(default_factory=dict)

    # Optional fields for specialized benchmarks
    system_prompt: str | None = None
    context: str | None = None
    constraints: list[str] | None = None
    test_cases: list[dict[str, Any]] | None = None

    @property
    def difficulty(self) -> DifficultyLevel:
        """Get difficulty level from metadata, or UNKNOWN if not set."""
        level = self.metadata.get("difficulty", "unknown")
        try:
            return DifficultyLevel(level)
        except ValueError:
            return DifficultyLevel.UNKNOWN

    @property
    def domain(self) -> str:
        """Get domain from metadata, or 'unknown' if not set."""
        return self.metadata.get("domain", "unknown")

    @property
    def task(self) -> str:
        """Get task name from metadata, or 'unknown' if not set."""
        return self.metadata.get("task", "unknown")


@dataclass
class EvaluationScore:
    """Result of evaluating a single prediction.

    Attributes:
        correct: Whether the prediction is correct (for binary metrics).
        score: Numeric score (0-1 for most metrics).
        details: Additional scoring details (e.g., which test cases passed).
    """

    correct: bool
    score: float
    details: dict[str, Any] = field(default_factory=dict)


class Benchmark(ABC):
    """Abstract base class for benchmarks.

    Subclasses must implement:
    - load(): Yields BenchmarkExample instances
    - evaluate(): Scores a prediction against ground truth

    Optionally override:
    - get_difficulty(): Compute difficulty for an example
    - get_oracle_models(): Models known to perform well
    """

    # Class attributes to be set by subclasses
    name: str
    version: str
    domain: BenchmarkDomain
    metric_type: MetricType
    description: str = ""

    # Benchmarks in EB1 training data (for OOD analysis)
    TRAINING_BENCHMARKS: set[str] = {
        # BBEH suite
        "bbeh_boardgame_qa",
        "bbeh_boolean_expressions",
        "bbeh_buggy_tables",
        "bbeh_causal_understanding",
        "bbeh_disambiguation_qa",
        "bbeh_dyck_languages",
        "bbeh_geometric_shapes",
        "bbeh_hyperbaton",
        "bbeh_linguini",
        "bbeh_movie_recommendation",
        "bbeh_multistep_arithmetic",
        # Code and math
        "livecodebench",
        "mbpp",
        "gsm8k",
        # Reasoning
        "agieval",
        "arc",
        # Knowledge
        "mmlu",
        "mmlu_pro_chemistry",
        "mmlu_pro_law",
        "mmlu_pro_biology",
        "mmlu_pro_business",
        "mmlu_pro_psychology",
        "mmlu_pro_history",
        "mmlu_pro_other",
        "truthfulqa",
        # Commonsense
        "hellaswag",
        "winogrande",
        # Reading/instruction
        "squad",
        "multi_news",
        "alpaca_data",
    }

    @abstractmethod
    def load(self, subset: str | None = None) -> Iterator[BenchmarkExample]:
        """Load examples from the benchmark.

        Args:
            subset: Optional subset name (e.g., "test", "validation", "mini").
                If None, loads the default evaluation split.

        Yields:
            BenchmarkExample instances.

        Raises:
            FileNotFoundError: If benchmark data not found.
            ValueError: If subset is invalid.
        """

    @abstractmethod
    def evaluate(
        self,
        prediction: str,
        example: BenchmarkExample,
    ) -> EvaluationScore:
        """Score a prediction against the ground truth.

        Args:
            prediction: Model's output.
            example: The benchmark example with ground truth.

        Returns:
            EvaluationScore with correctness and numeric score.
        """

    def get_difficulty(self, example: BenchmarkExample) -> DifficultyLevel:
        """Compute or retrieve difficulty level for an example.

        Default implementation returns UNKNOWN. Subclasses can override
        to compute difficulty from ground truth or metadata.
        """
        return example.difficulty

    def get_oracle_models(self) -> list[str] | None:
        """Return models known to perform well on this benchmark.

        Returns None if no oracle information is available.
        """
        return None

    def is_in_distribution(self) -> bool:
        """Check if this benchmark overlaps with EB1 training data."""
        return self.name.lower() in self.TRAINING_BENCHMARKS

    def __len__(self) -> int:
        """Return number of examples. May require loading data."""
        return sum(1 for _ in self.load())

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}(name='{self.name}', domain={self.domain.value})"
