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

GPQA (Graduate-level Google-Proof Q&A) is a challenging multiple-choice
benchmark requiring expert-level reasoning. Fresh for EB1 evaluation.
"""

import re
from typing import Iterator

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

# GPQA subsets available
GPQA_SUBSETS = {
    "diamond": "gpqa_diamond",  # Hardest subset (~198 examples)
    "main": "gpqa_main",  # Full dataset (~448 examples)
    "extended": "gpqa_extended",  # Extended with more examples
}


def _extract_answer_letter(prediction: str) -> str | None:
    """Extract answer letter (A/B/C/D) from model output.

    Handles common response patterns:
    - "A" or "(A)"
    - "The answer is A"
    - "Answer: A"
    - "A. option text"
    """
    text = prediction.strip().upper()

    # Direct letter at start
    if text and text[0] in "ABCD":
        return text[0]

    # Pattern: "The answer is X" or "Answer: X"
    match = re.search(r"(?:answer(?:\s+is)?[:\s]*)\(?([ABCD])\)?", text, re.IGNORECASE)
    if match:
        return match.group(1).upper()

    # Pattern: "(X)" anywhere
    match = re.search(r"\(([ABCD])\)", text, re.IGNORECASE)
    if match:
        return match.group(1).upper()

    # Last resort: any standalone letter
    match = re.search(r"\b([ABCD])\b", text)
    if match:
        return match.group(1).upper()

    return None


@register_benchmark("gpqa_diamond")
class GPQABenchmark(Benchmark):
    """GPQA Graduate-Level Q&A benchmark.

    Wraps HuggingFace Idavidrein/gpqa dataset. Multiple-choice format
    with 4 options (A-D). Evaluation extracts answer letter from output.

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

    name = "gpqa_diamond"
    version = "1.0"
    domain = BenchmarkDomain.REASONING
    metric_type = MetricType.EXACT_MATCH
    description = "Graduate-level expert Q&A requiring deep reasoning"

    def __init__(
        self,
        subset: str = "diamond",
        max_examples: int | None = None,
    ):
        """Initialize GPQA benchmark.

        Args:
            subset: Which GPQA subset ("diamond", "main", "extended").
            max_examples: Maximum examples to load (None = all).
        """
        if subset not in GPQA_SUBSETS:
            available = list(GPQA_SUBSETS.keys())
            raise ValueError(f"Unknown GPQA subset: {subset}. Available: {available}")
        self._subset = subset
        self._config = GPQA_SUBSETS[subset]
        self._max_examples = max_examples
        # Update name based on subset
        self.name = f"gpqa_{subset}"

    def load(self, subset: str | None = None) -> Iterator[BenchmarkExample]:
        """Load GPQA examples via HuggingFace.

        Args:
            subset: Override subset if provided.

        Yields:
            BenchmarkExample instances with multiple-choice format.
        """
        config = GPQA_SUBSETS.get(subset or self._subset, self._config)

        # GPQA is on HuggingFace: Idavidrein/gpqa
        source = HuggingFaceSource(
            "Idavidrein/gpqa",
            split="train",  # GPQA only has train split
            config=config,
        )

        count = 0
        for batch in source.read_batches():
            for row in batch:
                if self._max_examples and count >= self._max_examples:
                    return

                # GPQA format: Correct Answer + Incorrect Answer 1/2/3
                # Correct answer is always first (A)
                choices = {
                    "A": str(row.get("Correct Answer", "")),
                    "B": str(row.get("Incorrect Answer 1", "")),
                    "C": str(row.get("Incorrect Answer 2", "")),
                    "D": str(row.get("Incorrect Answer 3", "")),
                }

                # The correct answer is always A in raw GPQA (Correct Answer field)
                # But we should track which letter maps to correct
                correct_answer = "A"  # In raw GPQA, Correct Answer is listed first

                question = str(row.get("Question", ""))

                # Format question with choices
                formatted_query = question
                if choices:
                    formatted_query += "\n\nOptions:\n"
                    for letter, text in sorted(choices.items()):
                        if text:
                            formatted_query += f"{letter}. {text}\n"

                yield BenchmarkExample(
                    id=f"gpqa_{self._subset}_{count}",
                    query=formatted_query,
                    reference=correct_answer,
                    metadata={
                        "domain": "reasoning",
                        "source": "gpqa",
                        "subset": self._subset,
                        "subdomain": row.get("Subdomain", row.get("subdomain", "unknown")),
                        "choices": choices,
                        "raw_question": question,
                    },
                )
                count += 1

    def evaluate(
        self,
        prediction: str,
        example: BenchmarkExample,
    ) -> EvaluationScore:
        """Evaluate prediction by extracting answer letter.

        Args:
            prediction: Model's output (should contain A/B/C/D).
            example: Ground truth example.

        Returns:
            EvaluationScore with letter match result.
        """
        extracted = _extract_answer_letter(prediction)
        if isinstance(example.reference, str):
            reference = example.reference
        else:
            reference = example.reference[0]

        correct = extracted == reference.upper() if extracted else False

        return EvaluationScore(
            correct=correct,
            score=1.0 if correct else 0.0,
            details={
                "metric": "letter_match",
                "extracted_answer": extracted,
                "expected_answer": reference,
            },
        )
