"""CodeXGLUE Defect Detection dataset integration via the Hugging Face Hub.

This dataset contains 27,318 C functions labeled as vulnerable (defective) or
secure. The task is binary classification: given a C function, predict whether
it contains a security vulnerability. Evaluation uses exact match on "1"/"0".

HuggingFace: google/code_x_glue_cc_defect_detection
"""

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 = "codexglue_defect"

entry = get_entry(_DATASET_NAME)


def _normalize(row: Dict[str, object]) -> Dict[str, object]:
    """Normalize a CodeXGLUE defect detection row."""
    func = str(row.get("func") or "")
    # target is boolean in the HF dataset; normalize to "1"/"0" string.
    target_raw = row.get("target")
    if isinstance(target_raw, bool):
        target = "1" if target_raw else "0"
    else:
        target = str(int(bool(target_raw)))

    question = (
        "Is the following C function vulnerable to a security defect? "
        "Answer 1 for yes (defective), 0 for no (secure).\n\n"
        f"{func}"
    )
    metadata: Dict[str, object] = {
        "source_id": row.get("id") or row.get("commit_id"),
        "project": row.get("project"),
        "license": entry.license,
    }
    return {
        "question": question,
        "answer": target,
        "choices": {"0": "Not defective", "1": "Defective"},
        "metadata": {k: v for k, v in metadata.items() if v},
    }


def _build_source() -> DataTransformer:
    hf = entry.huggingface
    if hf is None:
        raise RuntimeError("CodeXGLUE defect 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="CodeXGLUE C function defect detection (binary classification)",
    size_bytes=0,
    example_count=27318,
    example_item={
        "question": "Is the following C function vulnerable...",
        "answer": "0",
        "choices": {"0": "Not defective", "1": "Defective"},
        "metadata": {"project": "qemu"},
    },
)


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


__all__ = ["register", "METADATA"]
