"""Answer checking for generated traces: math via math-verify, MCQ via boxed letter."""

import re

from math_verify import parse, verify


def last_boxed(text):
    """Return the contents of the last \\boxed{...} (brace-balanced), or None."""
    start = text.rfind("\\boxed")
    if start < 0:
        return None
    i = text.find("{", start)
    if i < 0:
        return None
    depth = 0
    for j in range(i, len(text)):
        if text[j] == "{":
            depth += 1
        elif text[j] == "}":
            depth -= 1
            if depth == 0:
                return text[i + 1 : j]
    return None


def check(content, gold, answer_type):
    """Return (is_correct, extracted_prediction)."""
    boxed = last_boxed(content or "")
    if answer_type == "mcq":
        if boxed is None:
            return False, None
        m = re.search(r"\b([A-J])\b", re.sub(r"\\text\{([^}]*)\}", r"\1", boxed))
        pred = m.group(1) if m else boxed.strip()
        return pred == gold.strip(), pred
    if boxed is None:
        return False, None
    plain = re.sub(r"\\text\{([^}]*)\}|\\textbf\{([^}]*)\}", r"\1\2", boxed).strip().lower()
    if gold.strip().lower() in ("yes", "no", "true", "false"):
        first = re.match(r"[a-z]+", plain)
        return bool(first) and first.group(0) == gold.strip().lower(), boxed
    try:
        gold_p = parse(f"${gold}$")
        pred_p = parse(f"$\\boxed{{{boxed}}}$")
        ok = verify(gold_p, pred_p)
    except Exception:
        ok = False
    if not ok:
        # Plain-text answers (Yes/No, words) that math-verify can't parse.
        norm = lambda s: re.sub(r"[\s\\{}$.]|\\text|\\mathrm", "", s).lower()
        ok = norm(boxed) == norm(gold)
    return bool(ok), boxed