"""Behavioral tests for exact sizes, restart safety and logical plan validation."""
import asyncio
import copy
import json
import tempfile
import unittest
from pathlib import Path


class AssemblyTests(unittest.TestCase):
    def test_preinstalled_parts_do_not_depend_on_unfinished_assembly(self):
        from assembly_dataset.world import make_episode, RECIPES
        from assembly_dataset.config import Config
        for index in range(100):
            episode = make_episode(index, Config())
            scene = episode['scene']
            last = len(RECIPES[scene['domain']][1]) - 1
            for component in scene['components']:
                if f"{component['id']}:stage:{last}" in scene['initial_state']:
                    self.assertEqual(component['dependencies'], [])

    def test_exact_counts_resume_and_split_grouping(self):
        from assembly_dataset.config import Config
        from assembly_dataset.pipeline import generate, export
        for count in (0, 1, 2, 7, 100):
            with self.subTest(count=count), tempfile.TemporaryDirectory() as tmp:
                config = Config(output_dir=Path(tmp), target_examples=count, shard_size=11)
                asyncio.run(generate(config))
                export(config)
                rows = [json.loads(line) for p in Path(tmp).glob('sft/*/*/[0-9][0-9][0-9][0-9][0-9][0-9].jsonl')
                        for line in p.read_text().splitlines()]
                self.assertEqual(len(rows), count)
                self.assertEqual(len({r['record_id'] for r in rows}), count)
                groups = {}
                for r in rows:
                    self.assertEqual(groups.setdefault(r['scenario_id'], r['split']), r['split'])
                asyncio.run(generate(config))
                export(config)
                self.assertEqual(sum(len(p.read_text().splitlines()) for p in Path(tmp).glob('sft/*/*/[0-9][0-9][0-9][0-9][0-9][0-9].jsonl')), count)
                config.target_examples = count + 4
                asyncio.run(generate(config))
                export(config)
                self.assertEqual(sum(len(p.read_text().splitlines()) for p in Path(tmp).glob('sft/*/*/[0-9][0-9][0-9][0-9][0-9][0-9].jsonl')), count + 4)

    def test_rejects_changed_generation_settings_on_resume(self):
        from assembly_dataset.config import Config
        from assembly_dataset.pipeline import generate
        with tempfile.TemporaryDirectory() as tmp:
            config = Config(output_dir=Path(tmp), target_examples=4)
            asyncio.run(generate(config))
            config.seed += 1
            with self.assertRaisesRegex(ValueError, 'config|конфиг'):
                asyncio.run(generate(config))

    def test_plans_and_branches_are_executable(self):
        from assembly_dataset.world import make_episode, validate_plan, validate_branch
        from assembly_dataset.config import Config
        for index in range(100):
            episode = make_episode(index, Config())
            validate_plan(episode, episode['plan'])
            validate_branch(episode, episode['branch'])
            if episode['plan']['actions']:
                invalid = copy.deepcopy(episode['plan'])
                invalid['actions'][0]['object'] = 'does_not_exist'
                with self.assertRaises(ValueError):
                    validate_plan(episode, invalid)
                invalid_branch = copy.deepcopy(episode['branch'])
                invalid_branch['nodes'][0]['on_failure'] = 'missing_node'
                with self.assertRaises(ValueError):
                    validate_branch(episode, invalid_branch)

    def test_missing_goal_and_cycles_are_rejected(self):
        from assembly_dataset.world import make_episode, validate_plan, validate_branch
        from assembly_dataset.config import Config
        episode = next(e for i in range(100) if (e := make_episode(i, Config()))['plan']['status'] == 'ready')
        broken = copy.deepcopy(episode['plan'])
        broken['actions'] = []
        with self.assertRaises(ValueError):
            validate_plan(episode, broken)
        branch = copy.deepcopy(episode['branch'])
        branch['nodes'][0]['on_success'] = branch['nodes'][0]['id']
        with self.assertRaises(ValueError):
            validate_branch(episode, branch)

    def test_failure_branch_requires_fresh_scene_before_another_action(self):
        from assembly_dataset.world import make_episode, validate_branch
        from assembly_dataset.config import Config
        episode = make_episode(0, Config(blocked_fraction=0))
        branch = copy.deepcopy(episode['branch'])
        retry = copy.deepcopy(branch['nodes'][0])
        retry['id'] = 'retry'
        branch['nodes'][0]['on_failure'] = 'retry'
        branch['nodes'].append(retry)
        with self.assertRaises(ValueError):
            validate_branch(episode, branch)

    def test_reproducibility_and_structural_diversity(self):
        from assembly_dataset.world import make_episode
        from assembly_dataset.config import Config
        config = Config()
        signatures = {make_episode(i, config)['signature'] for i in range(1000)}
        self.assertGreater(len(signatures), 990)
        self.assertEqual(make_episode(12, config), make_episode(12, config))

    def test_config_rejects_invalid_counts_modes_and_splits(self):
        from assembly_dataset.config import Config
        for kwargs in ({'target_examples': -1}, {'target_examples': 1.5}, {'modes': []},
                       {'modes': ['full', 'full']}, {'split': [0.8, 0.5, 0.1]}):
            with self.subTest(kwargs=kwargs), self.assertRaises(ValueError):
                Config(**kwargs)


if __name__ == '__main__':
    unittest.main()