"""SimpleQA benchmark loader wrapping ember.api.data.

SimpleQA is a factoid QA benchmark with exact-match evaluation.
Fresh for EB1 evaluation (not in router training data).
"""

from typing import Iterator

from ember.api.data import stream
from embergraphrouter.evals.benchmarks.base import (
    Benchmark,
    BenchmarkDomain,
    BenchmarkExample,
    EvaluationScore,
    MetricType,
)
from embergraphrouter.evals.benchmarks.registry import register_benchmark


def _normalize_answer(text: str) -> str:
    """Normalize answer for comparison: lowercase, strip, collapse whitespace."""
    return " ".join(text.lower().strip().split())


def _exact_match(prediction: str, reference: str | list[str]) -> bool:
    """Check exact match after normalization."""
    pred_norm = _normalize_answer(prediction)
    if isinstance(reference, str):
        return pred_norm == _normalize_answer(reference)
    return any(pred_norm == _normalize_answer(ref) for ref in reference)


@register_benchmark("simpleqa")
class SimpleQABenchmark(Benchmark):
    """SimpleQA factoid question answering benchmark.

    Wraps ember.api.data.stream("simpleqa") to provide BenchmarkExample
    instances. Evaluation uses exact string matching after normalization.

    This benchmark is NOT in EB1 router training data, making it suitable
    for uncontaminated evaluation.
    """

    name = "simpleqa"
    version = "1.0"
    domain = BenchmarkDomain.KNOWLEDGE
    metric_type = MetricType.EXACT_MATCH
    description = "Factoid QA with exact-match evaluation from OpenAI SimpleQA"

    def __init__(self, max_examples: int | None = None):
        """Initialize SimpleQA benchmark.

        Args:
            max_examples: Maximum examples to load (None = all ~4326).
        """
        self._max_examples = max_examples

    def load(self, subset: str | None = None) -> Iterator[BenchmarkExample]:
        """Load SimpleQA examples via ember.api.data.

        Args:
            subset: Ignored for SimpleQA (single split).

        Yields:
            BenchmarkExample instances.
        """
        iterator = stream("simpleqa", normalize="record")
        if self._max_examples:
            iterator = iterator.limit(self._max_examples)

        for idx, record in enumerate(iterator):
            yield BenchmarkExample(
                id=f"simpleqa_{idx}",
                query=record.question.text,
                reference=record.answer.text,
                metadata={
                    "domain": "knowledge",
                    "source": "simpleqa",
                    **record.metadata,
                },
            )

    def evaluate(
        self,
        prediction: str,
        example: BenchmarkExample,
    ) -> EvaluationScore:
        """Evaluate prediction using exact match.

        Args:
            prediction: Model's output.
            example: Ground truth example.

        Returns:
            EvaluationScore with exact match result.
        """
        correct = _exact_match(prediction, example.reference)
        return EvaluationScore(
            correct=correct,
            score=1.0 if correct else 0.0,
            details={"metric": "exact_match"},
        )
