"""DROP benchmark loader for discrete reasoning.

DROP (Discrete Reasoning Over Paragraphs) requires numerical/logical reasoning
over text, such as counting, arithmetic, or sorting. It tests whether models
can perform multi-step reasoning to produce specific answers.
"""

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 _parse_number(text: str) -> float | None:
    """Extract numeric value from text.

    Handles:
    - Plain integers: "42"
    - Decimals: "3.14"
    - Negative: "-5"
    - Words: "one", "two", etc.
    - Percentages: "50%"
    - Ordinals: "first", "second", etc.
    """
    text = text.strip().lower()

    # Handle word numbers
    word_to_num = {
        "zero": 0, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5,
        "six": 6, "seven": 7, "eight": 8, "nine": 9, "ten": 10,
        "eleven": 11, "twelve": 12, "thirteen": 13, "fourteen": 14,
        "fifteen": 15, "sixteen": 16, "seventeen": 17, "eighteen": 18,
        "nineteen": 19, "twenty": 20, "thirty": 30, "forty": 40,
        "fifty": 50, "sixty": 60, "seventy": 70, "eighty": 80, "ninety": 90,
        "hundred": 100, "thousand": 1000, "million": 1000000,
        "first": 1, "second": 2, "third": 3, "fourth": 4, "fifth": 5,
        "sixth": 6, "seventh": 7, "eighth": 8, "ninth": 9, "tenth": 10,
    }

    if text in word_to_num:
        return float(word_to_num[text])

    # Remove percentage sign
    text = text.replace("%", "").strip()

    # Remove commas from numbers
    text = text.replace(",", "")

    # Try parsing as number
    try:
        return float(text)
    except ValueError:
        pass

    # Try extracting first number from text
    match = re.search(r"-?\d+(?:\.\d+)?", text)
    if match:
        return float(match.group())

    return None


def _extract_answer_from_prediction(prediction: str) -> str:
    """Extract the final answer from model prediction.

    DROP answers are typically short (numbers, spans, dates).
    Look for patterns like "The answer is X" or just extract the last line.
    """
    text = prediction.strip()

    # Pattern: "The answer is X" or "Answer: X"
    match = re.search(
        r"(?:the\s+)?answer\s*(?:is|:)\s*(.+?)(?:\.|$)",
        text,
        re.IGNORECASE,
    )
    if match:
        return match.group(1).strip()

    # Pattern: "Therefore, X" or "So, X"
    match = re.search(r"(?:therefore|so)[,:]?\s*(.+?)(?:\.|$)", text, re.IGNORECASE)
    if match:
        return match.group(1).strip()

    # Last line heuristic
    lines = [ln.strip() for ln in text.split("\n") if ln.strip()]
    if lines:
        return lines[-1]

    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)


def _compute_numeric_match(prediction: str, references: list[str]) -> bool:
    """Check if predicted number matches any reference number."""
    pred_num = _parse_number(prediction)
    if pred_num is None:
        return False

    for ref in references:
        ref_num = _parse_number(ref)
        if ref_num is not None and abs(pred_num - ref_num) < 1e-6:
            return True

    return False


@register_benchmark("drop")
class DROPBenchmark(Benchmark):
    """DROP Discrete Reasoning Over Paragraphs benchmark.

    Tests numerical and logical reasoning over text passages.
    Requires counting, arithmetic, comparisons, and span extraction.

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

    name = "drop"
    version = "1.0"
    domain = BenchmarkDomain.REASONING
    metric_type = MetricType.F1
    description = "Discrete reasoning requiring numerical and logical analysis"

    def __init__(
        self,
        max_examples: int | None = None,
        split: str = "validation",
    ):
        """Initialize DROP benchmark.

        Args:
            max_examples: Maximum examples to load (None = all).
            split: Dataset split ("train" or "validation").
        """
        self._max_examples = max_examples
        self._split = split

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

        Args:
            subset: Override split if provided.

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

        # DROP on HuggingFace: ucinlp/drop
        source = HuggingFaceSource(
            "ucinlp/drop",
            split=split,
        )

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

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

                # DROP has answers_spans structure
                answers_spans = row.get("answers_spans", {})
                spans = answers_spans.get("spans", [])
                types = answers_spans.get("types", [])

                # Collect all valid reference answers
                references = [str(s) for s in spans if s]
                if not references:
                    # Some examples use number format
                    number = row.get("answer", {}).get("number", "")
                    if number:
                        references = [str(number)]

                # Skip if no valid references
                if not references:
                    continue

                # Determine answer type
                answer_type = types[0] if types else "span"

                # Format query with passage
                formatted_query = f"Passage: {passage}\n\nQuestion: {question}"

                yield BenchmarkExample(
                    id=f"drop_{count}",
                    query=formatted_query,
                    reference=references,
                    context=passage,
                    metadata={
                        "domain": "reasoning",
                        "source": "drop",
                        "answer_type": answer_type,
                        "raw_question": question,
                        "n_references": len(references),
                    },
                )
                count += 1

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

        DROP uses both exact match and F1 for evaluation.

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

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

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

        # Compute metrics
        exact_match = _compute_exact_match(extracted, references)
        numeric_match = _compute_numeric_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 any match type succeeds
        correct = exact_match or numeric_match or max_f1 >= 0.8

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