"""Natural Questions benchmark loader for open-domain QA.

Natural Questions contains real Google search queries paired with Wikipedia
articles. Answers are spans extracted from the articles. Tests open-domain
question answering ability.
"""

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


def _normalize_answer(text: str) -> str:
    """Normalize answer text for comparison.

    - Lowercase
    - Strip whitespace
    - Remove articles (a, an, the)
    - Remove punctuation
    """
    text = text.lower().strip()

    # Remove articles
    text = re.sub(r"\b(a|an|the)\b", " ", text)

    # Remove punctuation
    text = re.sub(r"[^\w\s]", " ", text)

    # Collapse whitespace
    text = " ".join(text.split())

    return text


def _compute_f1(prediction: str, reference: str) -> float:
    """Compute token-level F1 score.

    Args:
        prediction: Predicted answer text.
        reference: Ground truth answer text.

    Returns:
        F1 score between 0 and 1.

    """
    pred_tokens = set(_normalize_answer(prediction).split())
    ref_tokens = set(_normalize_answer(reference).split())

    if not pred_tokens or not ref_tokens:
        return 1.0 if pred_tokens == ref_tokens else 0.0

    common = pred_tokens & ref_tokens

    if not common:
        return 0.0

    precision = len(common) / len(pred_tokens)
    recall = len(common) / len(ref_tokens)

    return 2 * precision * recall / (precision + recall)


def _compute_exact_match(prediction: str, references: list[str]) -> bool:
    """Check exact match against any reference answer."""
    pred_normalized = _normalize_answer(prediction)
    return any(_normalize_answer(ref) == pred_normalized for ref in references)


# Patterns for answer extraction
_ANSWER_PATTERNS = [
    r"(?:the\s+)?answer\s*(?:is|:)\s*[\"']?(.+?)[\"']?(?:\.|,|\n|$)",
    r"(?:it\s+(?:is|was)|that\s+(?:is|was))\s+[\"']?(.+?)[\"']?(?:\.|,|\n|$)",
    r"^(.+?)(?:\.|,)\s*(?:This|It|The|In|During)",
]
_BOLD_PATTERNS = [r"\*\*([^*]+)\*\*", r"\*([^*]+)\*", r"__([^_]+)__", r"_([^_]+)_"]
_DATE_PATTERN = r"\b(\d{4}|\w+\s+\d{1,2},?\s+\d{4}|\w+\s+\d{4})\b"
_FILLER_WORDS = frozenset(["the", "is", "was", "it", "this", "that", "because"])


def _score_answer_candidate(candidate: str) -> float:
    """Score a candidate answer for NQ-style short answer quality.

    Args:
        candidate: Extracted answer candidate.

    Returns:
        Score where higher indicates better NQ answer characteristics.

    """
    score = 0.0
    # Prefer shorter answers
    if len(candidate) < 30:
        score += 2.0
    elif len(candidate) < 60:
        score += 1.0
    # Boost dates
    if re.match(r"^\d{4}$|^\w+\s+\d{1,2},?\s+\d{4}$|^\w+\s+\d{4}$", candidate):
        score += 3.0
    # Boost proper nouns
    if candidate[0].isupper() and not candidate.isupper():
        score += 1.0
    # Boost numbers
    if re.search(r"\d", candidate):
        score += 0.5
    # Penalize filler words
    words = candidate.lower().split()
    filler_count = sum(1 for w in words if w in _FILLER_WORDS)
    score -= filler_count * 0.3
    return score


def _extract_candidates(text: str) -> list[str]:
    """Extract answer candidates from text using multiple strategies.

    Args:
        text: Model prediction text.

    Returns:
        List of candidate answer strings.

    """
    candidates: list[str] = []

    # Strategy 1: Explicit answer patterns
    for pattern in _ANSWER_PATTERNS:
        match = re.search(pattern, text, re.IGNORECASE)
        if match:
            candidate = match.group(1).strip()
            if 1 < len(candidate) < 100:
                candidates.append(candidate)

    # Strategy 2: Markdown emphasis
    for pattern in _BOLD_PATTERNS:
        for match in re.findall(pattern, text):
            candidate = match.strip()
            if 1 < len(candidate) < 100:
                candidates.append(candidate)

    # Strategy 3: Quoted text
    for match in re.findall(r'"([^"]+)"', text):
        candidate = match.strip()
        if 1 < len(candidate) < 100:
            candidates.append(candidate)

    # Strategy 4: Date patterns
    for match in re.findall(_DATE_PATTERN, text):
        if 3 < len(match) < 50:
            candidates.append(match)

    # Strategy 5: First sentence core
    if len(text) > 100:
        first_sentence = re.split(r"[.!?]+\s+", text)[0].strip()
        if 10 < len(first_sentence) < 150:
            core = re.sub(
                r"^(?:The\s+|It\s+(?:is|was)\s+|Yes,?\s+|No,?\s+)",
                "",
                first_sentence,
                flags=re.IGNORECASE,
            )
            if core and core != first_sentence:
                candidates.append(core.strip())

    # Strategy 6: Short text fallback
    if len(text) < 50:
        candidates.append(text)

    return candidates


def _extract_short_answer(prediction: str) -> str:
    """Extract the short answer from model prediction.

    Uses multiple extraction strategies and ranks candidates by NQ answer
    characteristics (short, factual, typically a name/date/number/entity).

    Args:
        prediction: Full model prediction text.

    Returns:
        Extracted short answer string.

    """
    text = prediction.strip()
    candidates = _extract_candidates(text)

    if not candidates:
        return text[:50].strip()

    candidates.sort(key=_score_answer_candidate, reverse=True)
    return candidates[0]


@register_benchmark("natural_questions")
class NaturalQuestionsBenchmark(Benchmark):
    """Natural Questions open-domain QA benchmark.

    Real Google search queries with Wikipedia-based answers.
    Tests factoid question answering and reading comprehension.

    This benchmark is NOT in EB1 router training data.
    """

    name = "natural_questions"
    version = "1.0"
    domain = BenchmarkDomain.KNOWLEDGE
    metric_type = MetricType.F1
    description = "Open-domain QA from real Google queries"

    def __init__(
        self,
        max_examples: int | None = None,
        split: str = "validation",
        short_answers_only: bool = True,
    ):
        """Initialize Natural Questions benchmark.

        Args:
            max_examples: Maximum examples to load (None = all).
            split: Dataset split ("train" or "validation").
            short_answers_only: If True, only include examples with short answers.
        """
        self._max_examples = max_examples
        self._split = split
        self._short_answers_only = short_answers_only

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

        Args:
            subset: Override split if provided.

        Yields:
            BenchmarkExample instances with question and context.
        """
        split = subset or self._split

        # Natural Questions on HuggingFace: google-research-datasets/natural_questions
        # Using the simplified version for easier processing
        source = HuggingFaceSource(
            "google-research-datasets/nq_open",
            split=split,
        )

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

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

                # NQ open has list of short answers
                answers = row.get("answer", [])
                if isinstance(answers, str):
                    answers = [answers]
                answers = [str(a) for a in answers if a]

                # Skip if no valid answers
                if not answers:
                    continue

                # Format query (NQ open doesn't include context)
                formatted_query = f"Question: {question}\n\nProvide a short, factual answer."

                yield BenchmarkExample(
                    id=f"nq_{count}",
                    query=formatted_query,
                    reference=answers,
                    metadata={
                        "domain": "knowledge",
                        "source": "natural_questions",
                        "raw_question": question,
                        "n_answers": len(answers),
                    },
                )
                count += 1

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

        Args:
            prediction: Model's output.
            example: Ground truth with reference answers.

        Returns:
            EvaluationScore with F1 and exact match metrics.
        """
        # Extract short answer from prediction
        extracted = _extract_short_answer(prediction)

        # Get references
        references = (
            example.reference if isinstance(example.reference, list) else [example.reference]
        )

        # Compute metrics
        exact_match = _compute_exact_match(extracted, references)

        # Compute F1 against best matching reference
        f1_scores = [_compute_f1(extracted, ref) for ref in references]
        max_f1 = max(f1_scores) if f1_scores else 0.0

        # Consider correct if exact match or high F1
        correct = exact_match or max_f1 >= 0.7

        return EvaluationScore(
            correct=correct,
            score=max_f1,
            details={
                "metric": "f1_em",
                "f1": max_f1,
                "exact_match": exact_match,
                "extracted_answer": extracted,
                "references": references,
            },
        )
