"""Central registry describing supported Ember datasets and subtasks."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Mapping, Optional, Sequence


@dataclass(frozen=True)
class HuggingFaceLink:
    """Descriptor for loading a dataset from the Hugging Face Hub."""

    repo_id: str
    split: str
    subset: Optional[str] = None
    trust_remote_code: bool = False


@dataclass(frozen=True)
class DatasetSubtask:
    """Metadata describing a child task derived from a parent corpus."""

    identifier: str
    description: str
    filters: Mapping[str, Sequence[str]]
    eval_type: Optional[str] = None


@dataclass(frozen=True)
class DatasetCatalogEntry:
    """Declarative description for a benchmark dataset.

    Attributes:
        eval_type: Registry name of the evaluator to use for this dataset
            (e.g. ``"multiple_choice"``, ``"exact_match"``).  ``None`` means
            the caller must decide.  Subtask entries may override this.
    """

    name: str
    license: str
    description: str
    huggingface: Optional[HuggingFaceLink] = None
    subtasks: Sequence[DatasetSubtask] = field(default_factory=tuple)
    eval_type: Optional[str] = None


_CATALOG: Mapping[str, DatasetCatalogEntry] = {
    "simpleqa": DatasetCatalogEntry(
        name="simpleqa",
        license="Apache-2.0",
        description="Google SimpleQA verified evaluation set",
        huggingface=HuggingFaceLink(repo_id="google/simpleqa-verified", split="eval"),
        eval_type="exact_match",
    ),
    "fever_v2": DatasetCatalogEntry(
        name="fever_v2",
        license="CC-BY-SA-3.0",
        description="FEVER v2.0 fact verification",
        huggingface=HuggingFaceLink(
            repo_id="fever",
            split="validation",
            trust_remote_code=True,
        ),
        eval_type="multiple_choice",
    ),
    "mathvista": DatasetCatalogEntry(
        name="mathvista",
        license="CC-BY-SA-4.0",
        description="MathVista testmini split",
        huggingface=HuggingFaceLink(repo_id="AI4Math/MathVista", split="testmini"),
        eval_type="multiple_choice",
    ),
    "livecodebench": DatasetCatalogEntry(
        name="livecodebench",
        license="MIT",
        description="LiveCodeBench execution tasks",
        huggingface=HuggingFaceLink(repo_id="livecodebench/execution", split="test"),
        eval_type="python_literal",
    ),
    "halueval": DatasetCatalogEntry(
        name="halueval",
        license="CC-BY-NC-2.0",
        description="HaluEval hallucination classification benchmark",
        huggingface=HuggingFaceLink(repo_id="flowaicom/HaluEval", split="test"),
        eval_type="yes_no",
    ),
    "bbeh": DatasetCatalogEntry(
        name="bbeh",
        license="Apache-2.0",
        description="BBEH benchmark corpus",
        huggingface=HuggingFaceLink(repo_id="BBEH/bbeh", split="train"),
        eval_type="bbeh",
        subtasks=(
            DatasetSubtask(
                identifier="word_sorting",
                description="BBEH word sorting task",
                filters={"task": ("word sorting",)},
                eval_type="token_list",
            ),
            DatasetSubtask(
                identifier="buggy_tables",
                description="BBEH buggy tables task",
                filters={"task": ("buggy tables",)},
                eval_type="json_value",
            ),
        ),
    ),
    "mmlu_pro": DatasetCatalogEntry(
        name="mmlu_pro",
        license="Apache-2.0",
        description="MMLU-Pro evaluation set",
        huggingface=HuggingFaceLink(repo_id="TIGER-Lab/MMLU-Pro", split="test"),
        eval_type="multiple_choice",
        subtasks=(
            DatasetSubtask(
                identifier="health",
                description="MMLU-Pro health discipline",
                filters={"category": ("health",)},
            ),
        ),
    ),
    "agieval": DatasetCatalogEntry(
        name="agieval",
        license="MIT",
        description="AGIEval graduate-level reasoning benchmark",
        huggingface=HuggingFaceLink(
            repo_id="zacharyxxxxcr/AGIEval",
            split="validation",
            trust_remote_code=True,
        ),
        eval_type="multiple_choice",
        subtasks=(
            DatasetSubtask("aqua-rat", "AGIEval AQuA-RAT subset", {}),
            DatasetSubtask("gaokao-biology", "AGIEval Gaokao biology subset", {}),
            DatasetSubtask("gaokao-chemistry", "AGIEval Gaokao chemistry subset", {}),
            DatasetSubtask("gaokao-chinese", "AGIEval Gaokao Chinese subset", {}),
            DatasetSubtask("gaokao-english", "AGIEval Gaokao English subset", {}),
            DatasetSubtask("gaokao-geography", "AGIEval Gaokao geography subset", {}),
            DatasetSubtask("gaokao-history", "AGIEval Gaokao history subset", {}),
            DatasetSubtask("gaokao-mathcloze", "AGIEval Gaokao math cloze subset", {}),
            DatasetSubtask("gaokao-mathqa", "AGIEval Gaokao math QA subset", {}),
            DatasetSubtask("gaokao-physics", "AGIEval Gaokao physics subset", {}),
            DatasetSubtask("jec-qa-ca", "AGIEval JEC QA civil service subset", {}),
            DatasetSubtask("jec-qa-kd", "AGIEval JEC QA knowledge subset", {}),
            DatasetSubtask("logiqa-en", "AGIEval LogiQA English subset", {}),
            DatasetSubtask("logiqa-zh", "AGIEval LogiQA Chinese subset", {}),
            DatasetSubtask("lsat-ar", "AGIEval LSAT analytical reasoning subset", {}),
            DatasetSubtask("lsat-lr", "AGIEval LSAT logical reasoning subset", {}),
            DatasetSubtask("lsat-rc", "AGIEval LSAT reading comprehension subset", {}),
            DatasetSubtask("math", "AGIEval graduate math subset", {}),
            DatasetSubtask("sat-en", "AGIEval SAT English subset", {}),
            DatasetSubtask("sat-en-without-passage", "AGIEval SAT English (no passage) subset", {}),
            DatasetSubtask("sat-math", "AGIEval SAT math subset", {}),
        ),
    ),
    "cruxeval": DatasetCatalogEntry(
        name="cruxeval",
        license="MIT",
        description="CRUXEval code reasoning benchmark (800 Python functions)",
        huggingface=HuggingFaceLink(repo_id="cruxeval-org/cruxeval", split="test"),
        eval_type="python_literal",
    ),
    "cruxeval_x": DatasetCatalogEntry(
        name="cruxeval_x",
        license="MIT",
        description="CRUXEval-X multilingual code output prediction (19 languages)",
        huggingface=HuggingFaceLink(
            repo_id="CRUXEVAL-X/cruxeval-x",
            split="test",
            trust_remote_code=True,
        ),
        eval_type="python_literal",
    ),
    "codexglue_defect": DatasetCatalogEntry(
        name="codexglue_defect",
        license="MIT",
        description="CodeXGLUE C function defect detection (binary classification)",
        huggingface=HuggingFaceLink(
            repo_id="google/code_x_glue_cc_defect_detection",
            split="test",
        ),
        eval_type="exact_match",
    ),
    "mbpp": DatasetCatalogEntry(
        name="mbpp",
        license="CC-BY-4.0",
        description="Mostly Basic Python Problems (MBPP)",
        huggingface=HuggingFaceLink(repo_id="mbpp", split="test"),
        eval_type="code_competition",
    ),
    "table_arithmetic": DatasetCatalogEntry(
        name="table_arithmetic",
        license="Generated",
        description="Synthetic table arithmetic prompts",
        eval_type="numeric",
    ),
    "table_bias": DatasetCatalogEntry(
        name="table_bias",
        license="Generated",
        description="Synthetic table bias prompts",
        eval_type="multiple_choice",
    ),
    "supergpqa_management": DatasetCatalogEntry(
        name="supergpqa_management",
        license="ODC-Attribution",
        description="SuperGPQA management multiple-choice questions",
        eval_type="multiple_choice",
    ),
    "commongenhardened": DatasetCatalogEntry(
        name="commongenhardened",
        license="MIT",
        description="CommonGen-Hard constrained text generation",
        eval_type="exact_match",
    ),
    "mathvqa": DatasetCatalogEntry(
        name="mathvqa",
        license="CC-BY-SA-4.0",
        description="MathVQA multimodal mathematical reasoning",
        eval_type="exact_match",
    ),
}


def get_entry(name: str) -> DatasetCatalogEntry:
    try:
        return _CATALOG[name]
    except KeyError as exc:
        raise KeyError(f"Unknown dataset '{name}'") from exc


def get_eval_type(name: str) -> Optional[str]:
    """Return the evaluator registry name for a dataset.

    Handles dotted subtask names (e.g. ``"bbeh.word_sorting"``).  When a
    subtask declares its own ``eval_type`` that value takes precedence over
    the parent entry.

    Returns:
        Registry name such as ``"multiple_choice"`` or ``"exact_match"``,
        or ``None`` when the dataset is unknown or has no declared type.
    """
    parent_name = name
    subtask_id: Optional[str] = None
    if "." in name:
        parent_name, subtask_id = name.split(".", 1)

    entry = _CATALOG.get(parent_name)
    if entry is None:
        return None

    if subtask_id is not None:
        for subtask in entry.subtasks:
            if subtask.identifier == subtask_id and subtask.eval_type:
                return subtask.eval_type

    return entry.eval_type


def list_catalog_entries() -> Sequence[DatasetCatalogEntry]:
    return tuple(_CATALOG.values())


__all__ = [
    "DatasetCatalogEntry",
    "DatasetSubtask",
    "HuggingFaceLink",
    "get_entry",
    "get_eval_type",
    "list_catalog_entries",
]
