from __future__ import annotations

import json
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable


BOXED_PREFIX = r"\boxed{"


@dataclass(frozen=True)
class Problem:
    problem_id: int
    prompt: str
    answer: str
    available: bool
    resource_refs: tuple[str, ...] = ()
    image_refs: tuple[str, ...] = ()
    source_url: str = ""

    @classmethod
    def from_dict(cls, row: dict[str, Any]) -> "Problem":
        return cls(
            problem_id=int(row["id"]),
            prompt=str(row["prompt"]),
            answer=str(row["answer"]),
            available=bool(row.get("available", True)),
            resource_refs=tuple(str(value) for value in row.get("resource_refs", [])),
            image_refs=tuple(str(value) for value in row.get("image_refs", [])),
            source_url=str(row.get("source_url", "")),
        )

    def to_dict(self) -> dict[str, Any]:
        return {
            "id": self.problem_id,
            "prompt": self.prompt,
            "answer": self.answer,
            "available": self.available,
            "resource_refs": list(self.resource_refs),
            "image_refs": list(self.image_refs),
            "source_url": self.source_url,
        }


def load_jsonl(path: Path) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    with path.open(encoding="utf-8") as handle:
        for line_number, line in enumerate(handle, start=1):
            if not line.strip():
                continue
            try:
                value = json.loads(line)
            except json.JSONDecodeError as error:
                raise ValueError(f"{path}:{line_number}: invalid JSON: {error}") from error
            if not isinstance(value, dict):
                raise ValueError(f"{path}:{line_number}: expected a JSON object")
            rows.append(value)
    return rows


def write_jsonl(path: Path, rows: Iterable[dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n")


def load_problems(path: Path) -> list[Problem]:
    return [Problem.from_dict(row) for row in load_jsonl(path)]


def problem_map(path: Path) -> dict[int, Problem]:
    return {problem.problem_id: problem for problem in load_problems(path)}


def extract_boxed(text: str) -> str | None:
    """Return the contents of the last balanced ``\\boxed{...}``."""
    starts = [match.start() for match in re.finditer(re.escape(BOXED_PREFIX), text)]
    for start in reversed(starts):
        content_start = start + len(BOXED_PREFIX)
        depth = 1
        for index in range(content_start, len(text)):
            char = text[index]
            if char == "{":
                depth += 1
            elif char == "}":
                depth -= 1
                if depth == 0:
                    return text[content_start:index].strip()
    return None


def extract_final_answer(text: str) -> tuple[str | None, str]:
    boxed = extract_boxed(text)
    if boxed is not None:
        return boxed, "boxed"
    matches = re.findall(r"(?im)^\s*FINAL_ANSWER\s*:\s*(.+?)\s*$", text)
    if matches:
        return matches[-1].strip(), "final_answer"
    return None, "missing"


def normalize_answer(value: str | None) -> str | None:
    if value is None:
        return None
    normalized = value.strip().replace("−", "-").replace("–", "-")
    normalized = normalized.replace("\\,", "").replace(",", "")
    normalized = re.sub(r"\s+", "", normalized)
    if normalized.startswith("$") and normalized.endswith("$"):
        normalized = normalized[1:-1]
    return normalized.lower()


def answer_is_correct(candidate: str | None, expected: str) -> bool:
    return normalize_answer(candidate) == normalize_answer(expected)


def message_text(message: dict[str, Any]) -> str:
    content = message.get("content", "")
    if isinstance(content, str):
        return content
    if not isinstance(content, list):
        return ""
    parts: list[str] = []
    for part in content:
        if isinstance(part, str):
            parts.append(part)
        elif isinstance(part, dict):
            text = part.get("text") or part.get("content")
            if isinstance(text, str):
                parts.append(text)
    return "\n".join(parts)


def final_assistant_text(events: Iterable[dict[str, Any]]) -> str:
    answer = ""
    for event in events:
        if event.get("type") != "message_end":
            continue
        message = event.get("message")
        if isinstance(message, dict) and message.get("role") == "assistant":
            answer = message_text(message)
    return answer