"""CRUXEval dataset integration via the Hugging Face Hub.

CRUXEval (Code Reasoning, Understanding, and eXecution Evaluation) provides
800 Python functions with input-output pairs for two tasks:

- CRUXEval-O (output prediction): given function + input, predict output
- CRUXEval-I (input prediction): given function + output, predict input

Both tasks use exact-match evaluation and require no code execution.
"""

from __future__ import annotations

from typing import Callable, Dict, Optional

from ember.api.data import DatasetInfo, HuggingFaceSource
from ember.api.datasets.catalog import get_entry
from ember.api.datasets.utils import DataTransformer

_DATASET_NAME = "cruxeval"

entry = get_entry(_DATASET_NAME)


def _normalize_output(row: Dict[str, object]) -> Dict[str, object]:
    """Normalize a CRUXEval row for output prediction."""
    code = str(row.get("code") or "")
    inp = str(row.get("input") or "")
    out = str(row.get("output") or "")
    question = (
        f"Given the following Python function:\n\n{code}\n\n"
        f"What is the return value when called with: {inp}\n\n"
        f"Provide only the exact output value."
    )
    return {
        "question": question,
        "answer": out,
        "choices": {},
        "metadata": {
            "source_id": row.get("id"),
            "task": "output_prediction",
            "license": entry.license,
        },
    }


def _normalize_input(row: Dict[str, object]) -> Dict[str, object]:
    """Normalize a CRUXEval row for input prediction."""
    code = str(row.get("code") or "")
    inp = str(row.get("input") or "")
    out = str(row.get("output") or "")
    question = (
        f"Given the following Python function:\n\n{code}\n\n"
        f"The function returns: {out}\n\n"
        f"What input produces this output? Provide only the exact input value."
    )
    return {
        "question": question,
        "answer": inp,
        "choices": {},
        "metadata": {
            "source_id": row.get("id"),
            "task": "input_prediction",
            "license": entry.license,
        },
    }


def _build_source(
    transform_fn: Callable[[Dict[str, object]], Dict[str, object]],
) -> DataTransformer:
    hf = entry.huggingface
    if hf is None:
        raise RuntimeError("CRUXEval catalog entry missing Hugging Face link")
    source = HuggingFaceSource(
        hf.repo_id,
        split=hf.split,
        config=hf.subset,
        trust_remote_code=hf.trust_remote_code,
    )
    return DataTransformer(source, transform_fn=transform_fn)


_OUTPUT_META = DatasetInfo(
    name=f"{_DATASET_NAME}.output",
    description="CRUXEval output prediction (exact match)",
    size_bytes=0,
    example_count=800,
    example_item={
        "question": "Given the following Python function: ...",
        "answer": "[(4, 1), ...]",
        "choices": {},
        "metadata": {"task": "output_prediction"},
    },
)

_INPUT_META = DatasetInfo(
    name=f"{_DATASET_NAME}.input",
    description="CRUXEval input prediction (exact match)",
    size_bytes=0,
    example_count=800,
    example_item={
        "question": "Given the following Python function: ...",
        "answer": "[1, 1, 3, 1, 3, 1]",
        "choices": {},
        "metadata": {"task": "input_prediction"},
    },
)

METADATA = (
    DatasetInfo(
        name=_DATASET_NAME,
        description="CRUXEval code reasoning benchmark (800 Python functions)",
        size_bytes=0,
        example_count=800,
        example_item=_OUTPUT_META.example_item,
    ),
    _OUTPUT_META,
    _INPUT_META,
)


def register(
    register_fn: Callable[[str, DataTransformer, Optional[DatasetInfo]], None],
) -> None:
    register_fn(_DATASET_NAME, _build_source(_normalize_output), METADATA[0])
    register_fn(f"{_DATASET_NAME}.output", _build_source(_normalize_output), _OUTPUT_META)
    register_fn(f"{_DATASET_NAME}.input", _build_source(_normalize_input), _INPUT_META)


__all__ = ["register", "METADATA"]
