"""Answer normalization for exact-match scoring.

GAIA scoring is an exact string match, so the *content* of the answer matters
more than how the model phrased it. Two concerns are handled here, and kept as
separate, individually testable functions:

* :func:`extract_final_answer` — pull the answer out of any wrapper the model
  leaked despite instructions ("FINAL ANSWER: 42", "The answer is 42.",
  markdown bold, surrounding quotes).
* :func:`normalize_answer` — light, conservative cleanup of the extracted text
  (whitespace, trailing punctuation, thousands separators on bare numbers).

Normalization is deliberately conservative: over-aggressive rewriting can turn
a correct answer into a wrong one. When in doubt, leave the text alone.
"""

from __future__ import annotations

import re

# Labels the model might prepend even when told not to. Matched case-insensitively
# at the very start of the (stripped) text. Deliberately excludes words like
# "result" that are also plausible standalone answers.
_LABEL_PATTERN = re.compile(
    r"^\s*(?:final answer|final response|answer)\s*[:\-]?\s*",
    re.IGNORECASE,
)

# A leading conversational stem like "The answer is" / "It is".
_STEM_PATTERN = re.compile(
    r"^\s*(?:the answer is|the final answer is|it is|this is)\s+",
    re.IGNORECASE,
)

# A bare number that may carry thousands separators, e.g. "1,234,567".
_THOUSANDS_NUMBER = re.compile(r"^[+-]?\d{1,3}(?:,\d{3})+(?:\.\d+)?$")


def extract_final_answer(text: str) -> str:
    """Strip any leaked wrapper and return just the answer text.

    Handles, in order: trailing whitespace, a ``FINAL ANSWER:`` style label,
    a conversational stem, surrounding markdown bold/backticks, and surrounding
    quotes. If the model emitted multiple lines with a final-answer label on
    one of them, that line wins.
    """
    if text is None:
        return ""

    candidate = text.strip()
    if not candidate:
        return ""

    # Peel off any fenced/quoted wrapper around the whole response first, so a
    # multi-line ```...``` block collapses to its contents before line logic.
    candidate = _strip_wrappers(candidate)

    # If a "FINAL ANSWER:" label appears on its own line somewhere, prefer the
    # text after the last such label — that's the model's considered answer.
    # Require the captured answer to contain an alphanumeric character so the
    # regex can't backtrack into capturing the lone ":" separator.
    labelled = [
        m for m in re.finditer(
            r"(?im)^\s*final answer\s*[:\-]?\s*(?P<ans>.+?)\s*$", candidate
        )
        if re.search(r"\w", m.group("ans"))
    ]
    if labelled:
        candidate = labelled[-1].group("ans").strip()
    else:
        # Otherwise take the last non-empty line (models often reason first,
        # then state the answer last) and peel off any leading label/stem.
        lines = [ln.strip() for ln in candidate.splitlines() if ln.strip()]
        if lines:
            candidate = lines[-1]
        candidate = _LABEL_PATTERN.sub("", candidate, count=1)
        candidate = _STEM_PATTERN.sub("", candidate, count=1)

    return _strip_wrappers(candidate)


def normalize_answer(text: str) -> str:
    """Apply conservative, exact-match-friendly cleanup to an answer.

    Collapses internal whitespace, removes a single trailing sentence period
    (but never the decimal point inside a number), and drops thousands
    separators from bare numeric answers. Leaves everything else untouched.
    """
    if text is None:
        return ""

    cleaned = _strip_wrappers(text.strip())
    cleaned = re.sub(r"\s+", " ", cleaned).strip()

    # Drop a single trailing period if the result isn't a decimal number.
    if cleaned.endswith(".") and not re.search(r"\d\.\d*$", cleaned):
        cleaned = cleaned[:-1].rstrip()

    # A trailing period can leave wrappers exposed (e.g. "**Paris**." -> "**Paris**").
    cleaned = _strip_wrappers(cleaned)

    # "1,234" -> "1234" only when the whole answer is one such number.
    if _THOUSANDS_NUMBER.match(cleaned):
        cleaned = cleaned.replace(",", "")

    return cleaned


def finalize(raw_model_output: str) -> str:
    """Convenience: extract the answer, then normalize it. This is what the
    agent calls to turn a raw model response into a submittable answer."""
    return normalize_answer(extract_final_answer(raw_model_output))


def _strip_wrappers(text: str) -> str:
    """Remove surrounding markdown emphasis, code fences, and matching quotes.

    Applied repeatedly until the text stops changing, so nested wrappers like
    ``**`x`**`` or triple-backtick fences are fully removed.
    """
    s = text.strip()
    while True:
        before = s
        # Symmetric runs of backticks (inline code or ``` fences).
        m = re.match(r"^(`+)(.*?)(`+)$", s, re.DOTALL)
        if m and len(m.group(1)) == len(m.group(3)):
            s = m.group(2).strip()
        # Markdown bold/italic.
        for token in ("**", "*", "__", "_"):
            if len(s) >= 2 * len(token) and s.startswith(token) and s.endswith(token):
                s = s[len(token):-len(token)].strip()
                break
        # Matching surrounding quotes.
        if len(s) >= 2 and s[0] == s[-1] and s[0] in "\"'":
            s = s[1:-1].strip()
        if s == before:
            return s