"""tau2-bench (Sierra). Official release: github.com/sierra-research/tau2-bench.

We load the ModelScope mirror of the official task files
(evalscope/tau2-bench-data, same repo layout: tau2/domains/<domain>/tasks.json).
Each record is a full task: agent purpose + user scenario + evaluation criteria.
"""

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

_DOMAINS = ('airline', 'retail', 'telecom', 'mock')


@register_dataset(
    DatasetSpec(
        name='tau2_bench',
        source='HuggingFaceH4/tau2-bench-data',  # HF mirror of the official GitHub data
        subset='airline',  # or retail / telecom / mock; override with --subset
        split='test',
        task_type='agent',
        tags=['agent', 'tool_use', 'dialog'],
        description='tau2-bench agent-tool-dialog tasks (official content, HF mirror).',
        params={'hub': 'hf_raw', 'hf_files': ['domains/{subset}/tasks.json']},
    )
)
def tau2_bench():
    def to_sample(record: dict) -> Sample:
        desc = record.get('description') or {}
        scenario = record.get('user_scenario') or {}
        instructions = (scenario.get('instructions') or {}).get('task_instructions')
        return Sample(
            input=desc.get('purpose') or record.get('id', ''),
            target='',
            metadata={
                'id': record.get('id'),
                'notes': desc.get('notes'),
                'task_instructions': instructions,
                'user_scenario': scenario,
                'initial_state': record.get('initial_state'),
                'evaluation_criteria': record.get('evaluation_criteria'),
                'annotations': record.get('annotations'),
            },
        )

    return to_sample