"""BoolQ benchmark loader for yes/no question answering.

BoolQ (Boolean Questions) is a yes/no reading comprehension dataset.
Questions require understanding of a short passage to determine if
a statement is true or false.
"""

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 _extract_boolean_answer(prediction: str) -> bool | None:
    """Extract yes/no/true/false answer from model output.

    Handles various formats:
    - "Yes" / "No"
    - "True" / "False"
    - "The answer is yes"
    - Multi-line responses with conclusion
    """
    text = prediction.strip().lower()

    # Direct single word answer
    if text in ("yes", "true"):
        return True
    if text in ("no", "false"):
        return False

    # Pattern: starts with yes/no
    if text.startswith(("yes,", "yes.", "yes ")):
        return True
    if text.startswith(("no,", "no.", "no ")):
        return False

    # Pattern: "the answer is X"
    match = re.search(r"answer\s+is[:\s]*(yes|no|true|false)", text, re.IGNORECASE)
    if match:
        answer = match.group(1).lower()
        return answer in ("yes", "true")

    # Pattern: "therefore, X" or "so, X"
    match = re.search(r"(?:therefore|so)[,:\s]*(yes|no|true|false)", text, re.IGNORECASE)
    if match:
        answer = match.group(1).lower()
        return answer in ("yes", "true")

    # Pattern: ends with yes/no
    if text.endswith((" yes", " yes.", " true", " true.")):
        return True
    if text.endswith((" no", " no.", " false", " false.")):
        return False

    # Look for any yes/no in the last 50 characters
    last_part = text[-50:]
    yes_match = re.search(r"\b(yes|true)\b", last_part, re.IGNORECASE)
    no_match = re.search(r"\b(no|false)\b", last_part, re.IGNORECASE)

    if yes_match and not no_match:
        return True
    if no_match and not yes_match:
        return False

    # If both found, use the last one
    if yes_match and no_match:
        yes_pos = yes_match.start()
        no_pos = no_match.start()
        return yes_pos > no_pos

    return None


@register_benchmark("boolq")
class BoolQBenchmark(Benchmark):
    """BoolQ Boolean Questions benchmark.

    Yes/no reading comprehension requiring understanding of passages.
    Good for testing ceiling behavior (simple task for capable models).

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

    name = "boolq"
    version = "1.0"
    domain = BenchmarkDomain.KNOWLEDGE
    metric_type = MetricType.EXACT_MATCH
    description = "Boolean yes/no question answering over passages"

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

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

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

        Args:
            subset: Override split if provided.

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

        # BoolQ on HuggingFace: google/boolq (or boolq)
        source = HuggingFaceSource(
            "google/boolq",
            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", ""))
                passage = str(row.get("passage", ""))
                answer = row.get("answer", False)

                # Convert boolean to string
                reference = "yes" if answer else "no"

                # Format query with passage
                formatted_query = (
                    f"Passage: {passage}\n\n"
                    f"Question: {question}\n\n"
                    "Answer with 'yes' or 'no'."
                )

                yield BenchmarkExample(
                    id=f"boolq_{count}",
                    query=formatted_query,
                    reference=reference,
                    context=passage,
                    metadata={
                        "domain": "knowledge",
                        "source": "boolq",
                        "raw_question": question,
                        "answer_bool": answer,
                    },
                )
                count += 1

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

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

        Returns:
            EvaluationScore with boolean match result.
        """
        extracted = _extract_boolean_answer(prediction)

        # Get reference boolean
        if isinstance(example.reference, str):
            reference = example.reference
        else:
            reference = example.reference[0]
        expected = reference.lower() in ("yes", "true")

        # Handle extraction failure
        if extracted is None:
            return EvaluationScore(
                correct=False,
                score=0.0,
                details={
                    "metric": "boolean_match",
                    "extracted_answer": None,
                    "expected_answer": expected,
                    "parse_error": True,
                },
            )

        correct = extracted == expected

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