"""LongBench v2 (official source: THUDM/LongBench-v2)."""

from ..sample import Sample
from ..registry import register_dataset
from ..spec import DatasetSpec


@register_dataset(
    DatasetSpec(
        name='longbench_v2',
        source='THUDM/LongBench-v2',  # official: https://huggingface.co/datasets/THUDM/LongBench-v2
        split='train',  # the dataset ships a single split
        task_type='mcq',
        tags=['long_context'],
        description='LongBench v2 long-context MCQ (official). Context kept in metadata.',
    )
)
def longbench_v2():
    def to_sample(record: dict) -> Sample:
        return Sample(
            input=record['question'],
            choices=[record['choice_A'], record['choice_B'], record['choice_C'], record['choice_D']],
            target=str(record['answer']).strip(),
            metadata={
                'context': record['context'],  # the long document; eval-time prompt assembly
                'domain': record.get('domain'),
                'sub_domain': record.get('sub_domain'),
                'difficulty': record.get('difficulty'),
                'length': record.get('length'),
                'subset': record.get('length'),  # official subsets: short/medium/long
            },
        )

    return to_sample