import asyncio
import json
from pathlib import Path
import tempfile
import unittest

from assembly_dataset.config import Config
from assembly_dataset.pipeline import generate, export, validate
from assembly_dataset.world import make_episode, training_rows


class PairedExportTests(unittest.TestCase):
    def test_paired_shards_and_exact_count(self):
        with tempfile.TemporaryDirectory() as tmp:
            config = Config(output_dir=Path(tmp), target_examples=10, shard_size=2)
            asyncio.run(generate(config))
            report = export(config)
            self.assertEqual(report['examples'], 10)
            pairs = 0
            for path in Path(tmp).glob('sft/*/*/[0-9]*.jsonl'):
                if '.structured.' in path.name:
                    continue
                structured_path = path.with_suffix('.structured.jsonl')
                self.assertTrue(structured_path.exists())
                text_rows = list(map(json.loads, path.read_text().splitlines()))
                json_rows = list(map(json.loads, structured_path.read_text().splitlines()))
                self.assertEqual(len(text_rows), len(json_rows))
                for text, structured in zip(text_rows, json_rows):
                    self.assertEqual(text['record_id'], structured['record_id'])
                    self.assertEqual(text['messages'][1], structured['messages'][1])
                    self.assertNotEqual(text['messages'][0], structured['messages'][0])
                    answer = json.loads(structured['messages'][2]['content'])
                    self.assertEqual(set(answer), {'status', 'plan'})
                    self.assertNotRegex(str(answer), r'part_\d|slot_\d')
                    pairs += 1
            self.assertEqual(pairs, 10)
            self.assertEqual(validate(config)['examples_checked'], 10)

    def test_structured_skills_and_branches(self):
        config = Config(domains=['operator_assistance'], min_components=1, max_components=1,
                        blocked_fraction=0, step_sampling='all')
        episode = make_episode(0, config)
        rows = list(training_rows(episode, config, output_format='structured'))
        full = json.loads(rows[0]['messages'][2]['content'])
        actions = full['plan']
        self.assertEqual(actions[2], {'skill': 'pick', 'object': 'распределительная коробка №1'})
        self.assertEqual(actions[3]['target'], 'оператор')
        self.assertEqual(actions[6]['target'], 'зона хранения №1')
        self.assertNotIn('target', actions[7])
        for row in rows[1:-1]:
            self.assertEqual(len(json.loads(row['messages'][2]['content'])['plan']), 1)
        branch = json.loads(rows[-1]['messages'][2]['content'])
        for i, action in enumerate(branch['plan']):
            self.assertEqual(action['on_failure'], 'replan_required')
            self.assertEqual(action['on_success'], f'step_{i+2}' if i+1 < len(actions) else 'ready')

    def test_blocked_empty_plan(self):
        config = Config(domains=['operator_assistance'], min_components=1, max_components=1, blocked_fraction=1)
        for row in training_rows(make_episode(0, config), config, output_format='structured'):
            self.assertEqual(json.loads(row['messages'][2]['content']), {'status': 'blocked', 'plan': []})