from __future__ import annotations

import json
import os
import re
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any

EVAL_DIR = Path(__file__).resolve().parent.parent  # eval/
if str(EVAL_DIR) not in sys.path:
    sys.path.insert(0, str(EVAL_DIR))

from common.io import require_file  # noqa: E402

EN_BOXED_INSTR = "{question}\nPlease reason step by step, and put your final answer within \\boxed{{}}."


@dataclass(frozen=True)
class Problem:
    task_id: str
    question: str
    answer: str


# --------------------------------------------------------------------------- #
# Data
# --------------------------------------------------------------------------- #


def load(data_file: str | os.PathLike) -> list[Problem]:
    path = Path(data_file)
    require_file(path, "AIME 2026 download: https://huggingface.co/datasets/MathArena/aime_2026")

    problems: list[Problem] = []
    with open(path) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            row = json.loads(line)
            problems.append(
                Problem(
                    task_id=row["task_id"],
                    question=row["question"],
                    answer=str(row["answer"]),
                )
            )
    assert len({p.task_id for p in problems}) == len(problems), "duplicate task_id"
    return problems


# --------------------------------------------------------------------------- #
# prompt
# --------------------------------------------------------------------------- #


def build_question(problem: Problem) -> str:
    return EN_BOXED_INSTR.format(question=problem.question.strip())


def build_messages(problem: Problem) -> list[dict[str, Any]]:
    """Single-turn user message; rendered into the model chat template by the checkpoint encoder."""
    return [{"role": "user", "content": build_question(problem)}]


# --------------------------------------------------------------------------- #
# Answer extraction and scoring
# --------------------------------------------------------------------------- #


def extract_boxed(text: str) -> str | None:
    r"""Return the content of the last \boxed{...}, supporting nested braces (e.g. \boxed{\frac{1}{2}})."""
    token = "\\boxed"
    starts = [match.start() for match in re.finditer(re.escape(token), text)]
    for start in reversed(starts):
        open_brace = text.find("{", start + len(token))
        if open_brace < 0:
            continue
        depth = 0
        for pos in range(open_brace, len(text)):
            if text[pos] == "{":
                depth += 1
            elif text[pos] == "}":
                depth -= 1
                if depth == 0:
                    return text[open_brace + 1 : pos]
    return None


_NUMBER_RE = re.compile(r"(?<![\w.])(\d{1,3})(?![\w.])")


def normalize_answer(ans: str | None) -> str | None:
    """Normalize an answer into a comparable string: strip $, braces, whitespace, thousands separators and leading +/0."""
    if ans is None:
        return None
    text = ans.strip()
    text = text.strip("$").strip()
    text = text.replace("\\left", "").replace("\\right", "")
    text = text.replace("\\!", "").replace("\\,", "").replace("\\;", "")
    text = text.replace("{,}", ",") if "{,}" in text else text  # LaTeX thousands separator: \boxed{1{,}234}
    text = text.replace(",", "").replace(" ", "")
    text = text.rstrip(".")
    if text.startswith("+"):
        text = text[1:]
    if re.fullmatch(r"-?\d+", text):
        return str(int(text))
    return text


# The fallback scans only the tail: reasoning text holds many numbers, so scanning the whole
# text would easily pick up a value from an intermediate draft.
_TAIL_CHARS = 400


def extract_answer_with_source(text: str) -> tuple[str | None, str]:
    r"""Extract the answer and report which path produced it.

    1. `boxed`: the content of the last \boxed{...}. This is the format the prompt requests, so it
       is the most reliable path;
    2. `tail_number`: the last 0-999 integer after `</think>` (or within the final 400 characters).
       A fallback for malformed output, **higher risk**; score.py reports its share separately;
    3. `none`: nothing found.
    """
    boxed = extract_boxed(text)
    if boxed is not None:
        norm = normalize_answer(boxed)
        if norm:
            return norm, "boxed"
    tail = text.rsplit("</think>", 1)[-1] if "</think>" in text else text[-_TAIL_CHARS:]
    nums = _NUMBER_RE.findall(tail)
    if nums:
        return normalize_answer(nums[-1]), "tail_number"
    return None, "none"


def is_correct(pred: str | None, gold: str) -> bool:
    """Exact integer match; non-integer answers fall back to case- and whitespace-insensitive string equality."""
    if pred is None:
        return False
    pred_norm, gold_norm = normalize_answer(pred), normalize_answer(gold)
    if pred_norm is None or gold_norm is None:
        return False
    if re.fullmatch(r"-?\d+", pred_norm) and re.fullmatch(r"-?\d+", gold_norm):
        return int(pred_norm) == int(gold_norm)
    return pred_norm.lower() == gold_norm.lower()


# --------------------------------------------------------------------------- #
# Entry points used by run_eval.py / score.py
# --------------------------------------------------------------------------- #


def build_records(problems: list[Problem], *, limit: int | None = None) -> list[dict[str, Any]]:
    records = []
    for problem in problems:
        records.append(
            {
                "task_id": problem.task_id,
                "question": problem.question,
                "answer": problem.answer,
                "judge_question": build_question(problem),
                "messages": build_messages(problem),
            }
        )
        if limit is not None and len(records) >= limit:
            break
    return records


def score_rollout(rollout: dict[str, Any], record: dict[str, Any]) -> dict[str, Any]:
    """Score one rollout. The rollout must carry raw_text (the model output)."""
    pred, source = extract_answer_with_source(rollout.get("raw_text", ""))
    row = dict(rollout)
    row["pred"] = pred
    row["extract_source"] = source
    row["gold"] = record["answer"]
    row["correct"] = is_correct(pred, record["answer"])
    return row