"""Evaluator implementations used by Ember.

Includes text normalization primitives (``normalize_answer``,
``token_f1_score``, ``exact_match_score``), comprehensive BBEH
extraction (``extract_answer_bbeh``, ``evaluate_bbeh``), and
framework evaluator classes.
"""

from __future__ import annotations

import ast
import json
import math
import re
import subprocess
import sys
import unicodedata
from collections import Counter
from collections.abc import Callable, Mapping, Sequence
from typing import Generic, TypeVar

from .base_evaluator import EvaluationResult, IEvaluator
from .extractors import (
    FinalLetterExtractor,
    IOutputExtractor,
    RegexExtractor,
    extract_mc_answer,
)

T_out = TypeVar("T_out")
T_extracted = TypeVar("T_extracted")
T_truth = TypeVar("T_truth")


# ---------------------------------------------------------------------------
# Text normalization and scoring primitives
# ---------------------------------------------------------------------------


def normalize_answer(text: str) -> str:
    """Normalize text for evaluation using the SQuAD standard.

    Applies NFKD Unicode decomposition, lowercases, removes articles
    (a, an, the), strips all non-alphanumeric characters (including
    Unicode punctuation like smart quotes and em dashes), and collapses
    whitespace.

    Args:
        text: Raw text to normalize.

    Returns:
        Normalized string suitable for exact-match or F1 comparison.
    """
    result = unicodedata.normalize("NFKD", text)
    result = result.lower()
    result = re.sub(r"\b(a|an|the)\b", " ", result)
    result = re.sub(r"[^a-z0-9\s]", "", result)
    return " ".join(result.split())


def token_f1_score(
    prediction: str, ground_truth: str,
) -> tuple[float, float, float]:
    """Compute token-level F1 between prediction and ground truth.

    Both strings are passed through :func:`normalize_answer` before
    tokenization.  Special-cases yes/no/noanswer tokens: if one side
    is a boolean keyword and doesn't match the other, returns zero
    immediately (SQuAD convention).

    Named ``token_f1_score`` to avoid collision with classification-level
    F1 metrics.

    Args:
        prediction: Predicted text.
        ground_truth: Ground truth text.

    Returns:
        ``(f1, precision, recall)`` tuple.  All values in ``[0, 1]``.
    """
    norm_pred = normalize_answer(prediction)
    norm_gt = normalize_answer(ground_truth)

    zero = (0.0, 0.0, 0.0)

    if (
        norm_pred in ("yes", "no", "noanswer")
        and norm_pred != norm_gt
    ):
        return zero
    if (
        norm_gt in ("yes", "no", "noanswer")
        and norm_pred != norm_gt
    ):
        return zero

    pred_tokens = norm_pred.split()
    gt_tokens = norm_gt.split()
    common = Counter(pred_tokens) & Counter(gt_tokens)
    num_same = sum(common.values())

    if num_same == 0:
        return zero

    precision = num_same / len(pred_tokens)
    recall = num_same / len(gt_tokens)
    f1 = (2 * precision * recall) / (precision + recall)
    return (f1, precision, recall)


def exact_match_score(prediction: str, ground_truth: str) -> bool:
    """Check if prediction matches ground truth after normalization.

    Args:
        prediction: Predicted text.
        ground_truth: Ground truth text.

    Returns:
        ``True`` if :func:`normalize_answer` produces identical strings.
    """
    return normalize_answer(prediction) == normalize_answer(ground_truth)


class ComposedEvaluator(
    IEvaluator[T_out, T_truth],
    Generic[T_out, T_extracted, T_truth],
):
    """Apply an extractor before delegating to the base evaluator.

    When extraction yields an empty or falsy value, the result metadata
    includes ``extraction_failed=True`` so callers can distinguish genuine
    wrong answers from extraction failures.
    """

    def __init__(
        self,
        extractor: IOutputExtractor[T_out, T_extracted],
        base_evaluator: IEvaluator[T_extracted, T_truth],
    ) -> None:
        self.extractor = extractor
        self.base_evaluator = base_evaluator

    def evaluate(
        self, system_output: T_out, correct_answer: T_truth, **kwargs: object
    ) -> EvaluationResult:
        extracted_value = self.extractor.extract(system_output, **kwargs)
        result = self.base_evaluator.evaluate(extracted_value, correct_answer, **kwargs)
        if not extracted_value:
            metadata = dict(result.metadata) if result.metadata else {}
            metadata["extraction_failed"] = True
            return EvaluationResult(
                is_correct=result.is_correct,
                score=result.score,
                metadata=metadata,
            )
        return result


class ExactMatchEvaluator(IEvaluator[str, str]):
    """Check if two strings match after stripping whitespace."""

    def __init__(
        self,
        compare_fn: Callable[[str, str], bool] | None = None,
        *,
        case_sensitive: bool = True,
    ) -> None:
        self.case_sensitive = case_sensitive
        self.compare_fn = compare_fn or self._default_compare

    def _default_compare(self, str1: str, str2: str) -> bool:
        """Compare two strings with configurable case sensitivity."""

        left = str1.strip()
        right = str2.strip()
        if not self.case_sensitive:
            left = left.lower()
            right = right.lower()
        return left == right

    def evaluate(
        self, system_output: str, correct_answer: str, **kwargs: object
    ) -> EvaluationResult:
        is_correct = self.compare_fn(system_output, correct_answer)
        score = 1.0 if is_correct else 0.0
        return EvaluationResult(is_correct=is_correct, score=score)


class PythonLiteralEqualityEvaluator(IEvaluator[str, str]):
    """Compare values by parsing Python literals when possible."""

    def __init__(self, float_tol: float | None = None) -> None:
        self.float_tol = float_tol

    def _try_eval(self, value: str) -> tuple[bool, object]:
        text = value.strip()
        try:
            return True, ast.literal_eval(text)
        except (ValueError, SyntaxError):
            return False, text

    def _eq(self, left: object, right: object) -> bool:
        if isinstance(left, float) and isinstance(right, float) and self.float_tol is not None:
            return math.isclose(left, right, rel_tol=0.0, abs_tol=self.float_tol)
        if isinstance(left, (list, tuple)) and isinstance(right, (list, tuple)):
            if len(left) != len(right):
                return False
            return all(
                self._eq(l_item, r_item)
                for l_item, r_item in zip(left, right, strict=True)
            )
        if isinstance(left, dict) and isinstance(right, dict):
            if left.keys() != right.keys():
                return False
            return all(self._eq(left[key], right[key]) for key in left)
        return left == right

    def evaluate(
        self, system_output: str, correct_answer: str, **kwargs: object
    ) -> EvaluationResult:
        pred_parsed, pred_value = self._try_eval(system_output)
        gold_parsed, gold_value = self._try_eval(correct_answer)

        if pred_parsed and gold_parsed:
            is_correct = self._eq(pred_value, gold_value)
        else:
            pred_text = str(pred_value).strip()
            gold_text = str(gold_value).strip()
            is_correct = pred_text == gold_text

        return EvaluationResult(is_correct=is_correct, score=1.0 if is_correct else 0.0)


class NumericToleranceEvaluator(IEvaluator[float, float]):
    """Check whether a numeric output is within a tolerance of the expected value."""

    def __init__(self, tolerance: float = 0.01) -> None:
        self.tolerance = tolerance

    def evaluate(
        self, system_output: float, correct_answer: float, **kwargs: object
    ) -> EvaluationResult:
        difference = abs(system_output - correct_answer)
        rounded_diff = round(difference, 8)
        is_correct = rounded_diff <= self.tolerance
        base = abs(correct_answer) if correct_answer != 0 else 1.0
        score = max(0.0, 1.0 - rounded_diff / base)
        return EvaluationResult(is_correct=is_correct, score=score, metadata={"diff": rounded_diff})


class CodeExecutionEvaluator(IEvaluator[str, str]):
    """Execute Python code and compare stdout (unsafe; opt-in required).

    .. deprecated::
        Prefer :class:`~ember.utils.eval.code_execution.CodeCompetitionEvaluator`
        which provides a proper sandbox with resource limits, multi-language
        support, and security validation.  This evaluator runs code without
        any sandboxing.
    """

    def __init__(self, timeout: float = 5.0, *, allow_unsafe: bool = False) -> None:
        self.timeout = timeout
        self.allow_unsafe = allow_unsafe

    def evaluate(
        self, system_output: str, correct_answer: str, **kwargs: object
    ) -> EvaluationResult:
        if not self.allow_unsafe:
            raise RuntimeError(
                "CodeExecutionEvaluator is disabled by default because it executes arbitrary code. "
                "Pass allow_unsafe=True (trusted inputs only) or use a sandboxed evaluator."
            )

        process_result: subprocess.CompletedProcess[str]
        try:
            process_result = subprocess.run(
                args=[sys.executable, "-c", system_output],
                capture_output=True,
                text=True,
                timeout=self.timeout,
            )
        except subprocess.TimeoutExpired as timeout_error:
            return EvaluationResult(
                is_correct=False,
                score=0.0,
                metadata={"error": f"TimeoutExpired: {str(timeout_error)}"},
            )

        stdout_str = process_result.stdout.strip()
        expected_str = correct_answer.strip()
        is_correct = stdout_str == expected_str
        return EvaluationResult(
            is_correct=is_correct,
            score=1.0 if is_correct else 0.0,
            metadata={
                "stdout": process_result.stdout,
                "stderr": process_result.stderr,
                "exit_code": process_result.returncode,
            },
        )


class MultipleChoiceStrictEvaluator(IEvaluator[str, str]):
    """Robust multiple-choice evaluator that focuses on final answer letters."""

    def __init__(self, valid_letters: str = "ABCDEFGHIJ") -> None:
        self.valid_letters = tuple(ch for ch in valid_letters if ch.isalpha())
        self.extractor = FinalLetterExtractor(valid_letters="".join(self.valid_letters))

    def _map_choice_text(self, text: str, choices: Mapping[str, str] | None) -> str:
        if not choices:
            return ""
        normalized_map: dict[str, str] = {}
        for label, option_text in choices.items():
            normalized_map[option_text.strip().lower()] = label.strip().upper()
        lines = [line.strip().lower() for line in text.splitlines() if line.strip()]
        for candidate in reversed(lines):
            mapped = normalized_map.get(candidate)
            if mapped:
                return mapped
        candidate = text.strip().lower()
        return normalized_map.get(candidate, "")

    def evaluate(
        self, system_output: str, correct_answer: str, **kwargs: object
    ) -> EvaluationResult:
        expected = correct_answer.strip().upper()
        choices = kwargs.get("choices")
        extracted = self.extractor.extract(system_output)
        metadata: dict[str, object] = {"extracted": extracted}

        if not extracted and isinstance(choices, Mapping):
            mapped = self._map_choice_text(system_output, choices)
            if mapped:
                extracted = mapped
                metadata["mapped_from_text"] = True

        if not extracted:
            return EvaluationResult(is_correct=False, score=0.0, metadata=metadata)

        extracted = extracted.strip().upper()
        if extracted not in self.valid_letters:
            metadata["invalid_letter"] = extracted
            return EvaluationResult(is_correct=False, score=0.0, metadata=metadata)

        is_correct = extracted == expected
        return EvaluationResult(
            is_correct=is_correct,
            score=1.0 if is_correct else 0.0,
            metadata=metadata,
        )


class YesNoEvaluator(IEvaluator[str, str]):
    """Evaluator for binary yes/no style tasks with common synonym handling."""

    _ALIASES = {
        "yes": "yes",
        "y": "yes",
        "true": "yes",
        "correct": "yes",
        "pass": "yes",
        "no": "no",
        "n": "no",
        "false": "no",
        "fail": "no",
    }

    def evaluate(
        self, system_output: str, correct_answer: str, **kwargs: object
    ) -> EvaluationResult:
        def normalize(value: str) -> str:
            key = value.strip().lower()
            return self._ALIASES.get(key, key)

        pred = normalize(system_output)
        gold = normalize(correct_answer)
        is_correct = pred == gold and pred in {"yes", "no"}
        return EvaluationResult(is_correct=is_correct, score=1.0 if is_correct else 0.0)


def strip_latex_bbeh(response: str) -> str:
    """Strip LaTeX formatting from BBEH answers.

    Handles ``$...$``, ``\\boxed{}``, ``\\text{}``, ``\\texttt{}``.

    Args:
        response: Raw answer string potentially containing LaTeX.

    Returns:
        Answer with LaTeX formatting removed.
    """
    if response.startswith("$") and response.endswith("$"):
        response = response[1:-1]
    if "boxed{" in response and response.endswith("}"):
        response = response[:-1].split("boxed{")[1]
    if "text{" in response and response.endswith("}"):
        response = response[:-1].split("text{")[1]
    if "texttt{" in response and response.endswith("}"):
        response = response[:-1].split("texttt{")[1]
    return response


# Answer prefixes searched case-insensitively with rfind (last wins).
_BBEH_ANSWER_PREFIXES: tuple[str, ...] = (
    "the final answer is:",
    "the final answer is ",
    "the answer is:",
    "the answer is ",
    "final answer:",
    "final answer is:",
    "answer:",
)

_BBEH_MC_PATTERNS: tuple[re.Pattern[str], ...] = (
    re.compile(r"\*\*\(([A-Za-z])\)\*\*"),
    re.compile(
        r"(?:only|option|expression|answer|choice|statement)"
        r"\s+\**\(([A-Za-z])\)\**",
        re.IGNORECASE,
    ),
    re.compile(r"is\s+\**\(([A-Za-z])\)\**", re.IGNORECASE),
    re.compile(r"^\s*\(([A-Za-z])\)\s*$"),
)


def extract_answer_bbeh(sample: str) -> str:
    """Extract final answer from BBEH model response.

    Comprehensive extraction pipeline:
    1. Bold-strip, then case-insensitive ``rfind`` on answer prefixes
       (last occurrence wins over intermediate reasoning).
    2. Start-of-response yes/no/true/false/proved/disproved detection.
    3. ``\\boxed{}`` extraction.
    4. Bold standalone answer at end of response.
    5. End-of-prose yes/no/true/false/proved/disproved.
    6. Parenthesized MC letter patterns.
    7. Verbose numeric refinement (sentence with exactly one number).

    Args:
        sample: Full model response text.

    Returns:
        Extracted answer string, preprocessed and lowercased.
    """
    # Strip bold markdown for prefix matching.
    sample_clean = sample.replace("**", "")

    answer = sample_clean
    prefix_found = False
    sample_lower = sample_clean.lower()

    for answer_prefix in _BBEH_ANSWER_PREFIXES:
        if answer_prefix in sample_lower:
            idx = sample_lower.rfind(answer_prefix)
            answer = sample_clean[idx + len(answer_prefix):].strip()
            prefix_found = True
            break

    # Start-of-response yes/no/true/false/proved/disproved.
    if not prefix_found:
        start_match = re.match(
            r"^\s*(yes|no|true|false|proved|disproved)\s*[.,]",
            sample_clean, re.IGNORECASE,
        )
        if start_match:
            answer = start_match.group(1).lower()
            prefix_found = True

    # \boxed{} extraction.
    if not prefix_found:
        boxed_match = re.search(r"\\boxed\{([^}]+)\}", sample_clean)
        if boxed_match:
            answer = boxed_match.group(1).strip()
            prefix_found = True

    # Bold standalone answer at end: "Therefore **42**."
    if not prefix_found:
        bold_end = re.search(r"\*\*([^*]+?)\*\*[.!?]?\s*$", sample)
        if bold_end:
            candidate = bold_end.group(1).strip()
            if len(candidate) <= 40:
                answer = candidate
                prefix_found = True

    # End-of-prose yes/no/true/false/proved/disproved.
    if not prefix_found:
        end_yn = re.search(
            r"\b(yes|no|true|false|proved|disproved)\s*[.!?]?\s*$",
            sample_clean, re.IGNORECASE,
        )
        if end_yn:
            answer = end_yn.group(1).lower()
            prefix_found = True

    # MC letter patterns.
    if not prefix_found:
        for pattern in _BBEH_MC_PATTERNS:
            match = pattern.search(sample_clean, re.IGNORECASE)
            if match:
                answer = f"({match.group(1).upper()})"
                prefix_found = True
                break

    # Verbose numeric refinement.
    if prefix_found and len(answer.split()) > 3:
        numbers = re.findall(r"\b(\d+(?:\.\d+)?)\b", answer)
        if len(numbers) == 1:
            answer = numbers[0]

    # Clean up.
    if answer.endswith("."):
        answer = answer[:-1]
    answer = strip_latex_bbeh(answer)
    answer = answer.lower().strip()
    answer = answer.replace(", ", ",").replace("**", "")
    answer = answer.split("\n")[0]
    if answer.endswith("."):
        answer = answer[:-1]

    return answer


def fuzzy_match_bbeh(prediction: str, reference: str) -> bool:
    """Fuzzy match for BBEH evaluation.

    Handles parenthetical answers, numeric equivalence, quote
    normalization, bracket wrapping, and trailing punctuation.

    Args:
        prediction: Extracted and preprocessed prediction.
        reference: Ground truth answer (preprocessed).

    Returns:
        ``True`` if answers match under fuzzy rules.
    """
    if prediction == reference:
        return True

    # Parentheses: (a) vs a
    if len(prediction) == 3 and prediction[0] == "(" and prediction[-1] == ")":
        if prediction[1] == reference:
            return True
    if len(reference) == 3 and reference[0] == "(" and reference[-1] == ")":
        if reference[1] == prediction:
            return True

    # Numeric comparison.
    try:
        if float(prediction) == float(reference):
            return True
    except ValueError:
        pass

    # Quote normalization.
    if prediction.replace("'", "") == reference.replace("'", ""):
        return True

    # Bracket handling.
    if f"[{reference}]" == prediction or f"[{prediction}]" == reference:
        return True

    # Question mark tolerance.
    if prediction.endswith("?") and prediction[:-1] == reference:
        return True
    if reference.endswith("?") and reference[:-1] == prediction:
        return True

    return False


def evaluate_bbeh(prediction: str, ground_truth: str) -> float:
    """Evaluate BBEH task using official methodology with MC fallback.

    Extracts the final answer from the prediction and compares against
    ground truth using fuzzy matching.  For single-letter MC ground
    truths (e.g., ``"(e)"``, ``"a"``), falls back to
    :func:`~ember.utils.eval.extractors.extract_mc_answer` which
    handles diverse answer formats that the BBEH prefix extractor misses.

    Args:
        prediction: Full model response.
        ground_truth: Expected answer.

    Returns:
        ``1.0`` if correct, ``0.0`` otherwise.
    """
    extracted = extract_answer_bbeh(prediction)
    reference = ground_truth.strip().lower().replace(", ", ",")
    if fuzzy_match_bbeh(extracted, reference):
        return 1.0

    # MC fallback for subtasks like hyperbaton, geometric_shapes.
    ref_letter = reference.strip("() ")
    if len(ref_letter) == 1 and ref_letter.isalpha():
        mc_answer = extract_mc_answer(prediction)
        if mc_answer and mc_answer.lower() == ref_letter:
            return 1.0

    return 0.0


# Backward-compatible private aliases for existing callers.
_strip_latex_bbeh = strip_latex_bbeh
_extract_answer_bbeh = extract_answer_bbeh
_fuzzy_match_bbeh = fuzzy_match_bbeh

# Legacy helper aliases kept for test backward compatibility.
# These map to the comprehensive extraction pipeline above.
_BBEH_YES_NO_RE = re.compile(
    r"^\s*\**\s*(yes|no|true|false|proved|disproved)\s*\**\s*[.,]",
    re.IGNORECASE,
)


def _extract_by_prefix(text: str) -> str | None:
    """Try each known answer prefix; return text after the last match."""
    text_lower = text.replace("**", "").lower()
    for prefix in _BBEH_ANSWER_PREFIXES:
        if prefix in text_lower:
            idx = text_lower.rfind(prefix)
            return text.replace("**", "")[idx + len(prefix):].strip()
    return None


def _extract_yes_no_start(text: str) -> str | None:
    """Detect a leading yes/no/true/false/proved/disproved token."""
    match = _BBEH_YES_NO_RE.match(text)
    if match:
        return match.group(1).lower()
    return None


def _extract_mc_letter(text: str) -> str | None:
    """Extract a parenthesized multiple-choice letter."""
    for pattern in _BBEH_MC_PATTERNS:
        match = pattern.search(text)
        if match:
            return f"({match.group(1).upper()})"
    return None


def _clean_bbeh_answer(raw: str) -> str:
    """Normalize a raw BBEH answer for comparison."""
    answer = strip_latex_bbeh(raw)
    answer = answer.lower().strip()
    answer = answer.replace(", ", ",").replace("**", "")
    answer = answer.split("\n")[0]
    if answer.endswith("."):
        answer = answer[:-1]
    return answer


class BBEHEvaluator(IEvaluator[str, str]):
    """Evaluator for BBEH using official extraction and fuzzy matching.

    Uses :func:`evaluate_bbeh` which includes MC fallback for subtasks
    like hyperbaton and geometric_shapes.
    """

    def evaluate(
        self, system_output: str, correct_answer: str, **kwargs: object
    ) -> EvaluationResult:
        gold = correct_answer.strip()
        if gold.startswith("[") and gold.endswith("]"):
            raise ValueError(
                f"BBEH ground truth must be scalar, got list-like: {gold!r}"
            )

        extracted = extract_answer_bbeh(system_output)
        reference = gold.lower().replace(", ", ",")
        score = evaluate_bbeh(system_output, correct_answer)
        is_correct = score >= 1.0
        metadata: dict[str, str] = {
            "extracted": extracted, "reference": reference,
        }
        return EvaluationResult(
            is_correct=is_correct,
            score=score,
            metadata=metadata,
        )


class TokenListEvaluator(IEvaluator[str, Sequence[str]]):
    """Compare ordered token sequences extracted from text outputs."""

    def __init__(self, separator_pattern: str = r"[,\s]+") -> None:
        self._separator = re.compile(separator_pattern)

    def _to_tokens(self, value: object) -> list[str]:
        if isinstance(value, (list, tuple)):
            return [str(item).strip() for item in value]
        text = str(value).strip()
        if not text:
            return []
        return [token for token in self._separator.split(text) if token]

    def _final_nonempty_line(self, text: str) -> str:
        for line in reversed(text.splitlines()):
            stripped = line.strip()
            if stripped:
                return stripped
        return text.strip()

    def evaluate(
        self, system_output: str, correct_answer: Sequence[str] | str, **kwargs: object
    ) -> EvaluationResult:
        gold_tokens = self._to_tokens(correct_answer)
        prediction_line = self._final_nonempty_line(system_output)
        predicted_tokens = self._to_tokens(prediction_line)
        is_correct = predicted_tokens == gold_tokens
        metadata = {"pred_tokens": predicted_tokens, "gold_tokens": gold_tokens}
        return EvaluationResult(
            is_correct=is_correct,
            score=1.0 if is_correct else 0.0,
            metadata=metadata,
        )


class JsonValueEvaluator(IEvaluator[str, str]):
    """Compare JSON-encoded results focusing on a single key (default ``result``)."""

    def __init__(self, key: str = "result", float_tol: float | None = None) -> None:
        self.key = key
        self.float_tol = float_tol

    def _parse(self, value: object) -> tuple[bool, object]:
        if isinstance(value, Mapping):
            return True, value
        if isinstance(value, str):
            try:
                return True, json.loads(value)
            except json.JSONDecodeError:
                return False, value.strip()
        return False, value

    def _compare_scalar(self, left: object, right: object) -> bool:
        if isinstance(left, (int, float)) and isinstance(right, (int, float)):
            if self.float_tol is not None:
                return math.isclose(float(left), float(right), rel_tol=0.0, abs_tol=self.float_tol)
            return float(left) == float(right)
        return left == right

    def evaluate(
        self, system_output: str, correct_answer: str, **kwargs: object
    ) -> EvaluationResult:
        pred_ok, pred_value = self._parse(system_output)
        gold_ok, gold_value = self._parse(correct_answer)

        if (
            pred_ok
            and isinstance(pred_value, Mapping)
            and self.key in pred_value
            and gold_ok
            and isinstance(gold_value, Mapping)
            and self.key in gold_value
        ):
            pred_result = pred_value[self.key]
            gold_result = gold_value[self.key]
            is_correct = self._compare_scalar(pred_result, gold_result)
        else:
            is_correct = system_output.strip() == str(correct_answer).strip()

        return EvaluationResult(is_correct=is_correct, score=1.0 if is_correct else 0.0)


class PartialRegexEvaluator(ComposedEvaluator[str, str, str]):
    """Extract with a regex, then check exact match."""

    def __init__(self, pattern: str) -> None:
        extractor = RegexExtractor(pattern)
        evaluator = ExactMatchEvaluator()
        super().__init__(extractor=extractor, base_evaluator=evaluator)
