"""Framework-neutral access to the frozen MATH task catalog."""

from __future__ import annotations

from dataclasses import dataclass
import hashlib
from importlib.resources import files
import json
from typing import Any

_TASK_FILES = {"development": "dev.json", "test": "test.json"}
_TASK_FILE_SHA256 = {
    "dev.json": "cab1d3012159917dc10d506219538037c0865e77106fd3f441bde1d203266241",
    "test.json": "27dcf066ea6ac7949c0b77d08b14917f569fabfc678d8759517f29255b284858",
}
_REQUIRED_FIELDS = {
    "task_idx",
    "source_id",
    "source_split",
    "source_class",
    "level",
    "instruction",
    "solution",
    "answer",
    "source_group",
}


@dataclass(frozen=True, slots=True)
class MathTask:
    task_id: int
    source_id: str
    source_split: str
    source_class: str
    level: str
    instruction: str
    solution: str
    answer: str
    source_group: str


@dataclass(frozen=True, slots=True)
class MathSplit:
    name: str
    role: str
    source: str
    source_revision: str
    split_policy: str
    tasks: tuple[MathTask, ...]
    sha256: str


def _read_split(name: str, filename: str) -> MathSplit:
    raw = files("math_env").joinpath("assets", "tasks", filename).read_bytes()
    digest = hashlib.sha256(raw).hexdigest()
    if digest != _TASK_FILE_SHA256[filename]:
        raise ValueError(
            f"MATH task-file hash mismatch for {filename}: "
            f"{digest} != {_TASK_FILE_SHA256[filename]}"
        )
    payload: Any = json.loads(raw)
    expected_keys = {
        "schema_version",
        "benchmark",
        "source",
        "source_revision",
        "split_policy",
        "role",
        "tasks",
    }
    if not isinstance(payload, dict) or set(payload) != expected_keys:
        raise ValueError(f"invalid MATH split schema: {filename}")
    if payload["schema_version"] != 1 or payload["benchmark"] != "math":
        raise ValueError(f"invalid MATH split identity: {filename}")
    if payload["role"] != name:
        raise ValueError(f"invalid MATH split role: {filename}")

    tasks: list[MathTask] = []
    for index, record in enumerate(payload["tasks"]):
        if not isinstance(record, dict) or set(record) != _REQUIRED_FIELDS:
            raise ValueError(f"invalid MATH task {index}: {filename}")
        task = MathTask(
            task_id=int(record["task_idx"]),
            source_id=str(record["source_id"]),
            source_split=str(record["source_split"]),
            source_class=str(record["source_class"]),
            level=str(record["level"]),
            instruction=str(record["instruction"]),
            solution=str(record["solution"]),
            answer=str(record["answer"]),
            source_group=str(record["source_group"]),
        )
        if not task.instruction.strip() or not task.answer.strip():
            raise ValueError(f"invalid MATH task text: {task.task_id}")
        tasks.append(task)
    ids = [task.task_id for task in tasks]
    if len(ids) != len(set(ids)):
        raise ValueError(f"duplicate MATH task IDs: {filename}")
    return MathSplit(
        name=name,
        role=str(payload["role"]),
        source=str(payload["source"]),
        source_revision=str(payload["source_revision"]),
        split_policy=str(payload["split_policy"]),
        tasks=tuple(tasks),
        sha256=digest,
    )


_SPLITS = {name: _read_split(name, filename) for name, filename in _TASK_FILES.items()}
_TASKS = {task.task_id: task for split in _SPLITS.values() for task in split.tasks}
if len(_TASKS) != sum(len(split.tasks) for split in _SPLITS.values()):
    raise ValueError("MATH splits overlap")
if {task.source_group for task in _SPLITS["development"].tasks} & {
    task.source_group for task in _SPLITS["test"].tasks
}:
    raise ValueError("MATH source groups cross split boundaries")
if set(_TASKS) != set(range(2000)):
    raise ValueError("MATH global task IDs are not the frozen contiguous 0..1999 range")


def get_split(name: str) -> MathSplit:
    try:
        return _SPLITS[name]
    except KeyError as error:
        raise ValueError(f"unknown MATH split: {name!r}") from error


def get_task(task_id: int) -> MathTask:
    try:
        return _TASKS[int(task_id)]
    except KeyError as error:
        raise ValueError(f"unknown MATH task ID: {task_id!r}") from error


def catalog_sha256() -> str:
    rows = [
        [
            split.role,
            task.task_id,
            task.source_id,
            task.source_class,
            task.level,
            task.instruction,
            task.solution,
            task.answer,
            task.source_group,
        ]
        for split in (_SPLITS["development"], _SPLITS["test"])
        for task in split.tasks
    ]
    canonical = json.dumps(rows, separators=(",", ":"), ensure_ascii=False)
    return hashlib.sha256(canonical.encode("utf-8")).hexdigest()