"""CRUXEval-X multilingual code reasoning dataset integration.

CRUXEval-X extends CRUXEval to 19 programming languages with content-consistent
tests. Each language has a minimum of 600 output-prediction problems, totaling
approximately 19,000 test cases. Evaluation uses exact match on predicted output.

Source: https://github.com/CRUXEVAL-X/cruxeval-x
HuggingFace: CRUXEVAL-X/cruxeval-x
"""

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_x"

entry = get_entry(_DATASET_NAME)

# Languages available in CRUXEval-X.
_LANGUAGES = (
    "C", "C++", "C#", "D", "Dart", "Go", "Haskell", "Java", "JavaScript",
    "Julia", "Kotlin", "Lua", "PHP", "Perl", "Python", "R", "Ruby", "Rust",
    "Scala",
)


def _normalize(row: Dict[str, object]) -> Dict[str, object]:
    """Normalize a CRUXEval-X row for output prediction."""
    code = str(row.get("code") or "")
    inp = str(row.get("input") or "")
    out = str(row.get("output") or "")
    lang = str(row.get("language") or row.get("lang") or "")
    question = (
        f"Given the following {lang} 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."
    )
    metadata: Dict[str, object] = {
        "source_id": row.get("id"),
        "language": lang,
        "license": entry.license,
    }
    return {
        "question": question,
        "answer": out,
        "choices": {},
        "metadata": {k: v for k, v in metadata.items() if v},
    }


def _build_source() -> DataTransformer:
    hf = entry.huggingface
    if hf is None:
        raise RuntimeError("CRUXEval-X 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=_normalize)


METADATA = DatasetInfo(
    name=_DATASET_NAME,
    description="CRUXEval-X multilingual code output prediction (19 languages)",
    size_bytes=0,
    example_count=19000,
    example_item={
        "question": "Given the following Python function: ...",
        "answer": "[1, 2, 3]",
        "choices": {},
        "metadata": {"language": "Python"},
    },
)


def register(
    register_fn: Callable[[str, DataTransformer, Optional[DatasetInfo]], None],
) -> None:
    register_fn(_DATASET_NAME, _build_source(), METADATA)


__all__ = ["register", "METADATA"]
