"""Grading logic for the independent AIME 2025 / HumanEval eval."""

import re
import subprocess
import sys
import tempfile

BOXED_RE = re.compile(r"\\boxed\{([^{}]*)\}")
FENCE_RE = re.compile(r"```(?:python)?\s*\n(.*?)```", re.DOTALL)
INT_RE = re.compile(r"-?\d[\d,]*")


def extract_aime_answer(text):
    """Return the model's integer answer: last \\boxed{} if present, else last integer."""
    if not text:
        return None
    boxed = BOXED_RE.findall(text)
    candidates = boxed if boxed else INT_RE.findall(text)
    for cand in reversed(candidates):
        digits = cand.strip().replace(",", "")
        try:
            return int(digits)
        except ValueError:
            continue
    return None


def grade_aime(text, expected):
    got = extract_aime_answer(text)
    return got is not None and got == int(expected)


ANSWER_LETTER_RE = re.compile(
    r"(?:answer|correct option|choice)\b[^A-D]{0,20}\b([A-D])\b", re.IGNORECASE
)
STANDALONE_LETTER_RE = re.compile(r"\(?\b([A-D])\b\)?")


def extract_mcq_letter(text):
    """Return the model's A-D choice: prefer an 'Answer: X' style match, else the
    last standalone capital letter A-D."""
    if not text:
        return None
    answered = ANSWER_LETTER_RE.findall(text)
    if answered:
        return answered[-1].upper()
    standalone = [m for m in STANDALONE_LETTER_RE.findall(text) if m.isupper()]
    return standalone[-1] if standalone else None


def extract_code(text):
    """Return the last fenced code block, or the raw text if no fence exists."""
    if not text:
        return ""
    blocks = FENCE_RE.findall(text)
    return blocks[-1].strip() if blocks else text.strip()


def run_humaneval_check(code, test_src, entry_point, timeout=15):
    """Execute candidate code against the HumanEval test in a subprocess.

    Returns (passed: bool, detail: str).
    """
    program = "\n".join(
        [
            "from typing import *",
            "import math",
            "",
            code,
            "",
            test_src,
            f"check({entry_point})",
            "print('HUMANEVAL_OK')",
        ]
    )
    with tempfile.TemporaryDirectory() as sandbox:
        try:
            proc = subprocess.run(
                [sys.executable, "-c", program],
                cwd=sandbox,
                capture_output=True,
                text=True,
                timeout=timeout,
            )
        except subprocess.TimeoutExpired:
            return False, "timeout"
    if proc.returncode == 0 and "HUMANEVAL_OK" in proc.stdout:
        return True, "ok"
    return False, (proc.stderr or proc.stdout).strip()[-500:]