"""OpenAI MRCR (multi-round coreference, long-context). Mirror: openai-mirror/mrcr."""

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


@register_dataset(
    DatasetSpec(
        name='openai_mrcr',
        source='openai-mirror/mrcr',  # OpenAI MRCR mirror on ModelScope
        subset='2needle',  # or 4needle / 8needle
        split='test',
        task_type='qa',
        tags=['long_context'],
        description='OpenAI MRCR long-context retrieval/coreference (community mirror).',
        params={'hub': 'modelscope'},
    )
)
def openai_mrcr():
    def to_sample(record: dict) -> Sample:
        return Sample(
            input=record['prompt'],
            target=str(record['answer']).strip(),
            metadata={
                'n_needles': record.get('n_needles'),
                'total_messages': record.get('total_messages'),
                'random_string_to_prepend': record.get('random_string_to_prepend'),
            },
        )

    return to_sample