import asyncio
import json
from pathlib import Path
import tempfile
from types import SimpleNamespace
import unittest
from unittest.mock import patch

from assembly_dataset.config import Config, TeacherConfig
from assembly_dataset.pipeline import generate, export, validate
from assembly_dataset.teacher import Plan
from assembly_dataset.world import solve
from assembly_dataset.language import action_text
from llm.config import OpenAiConfig


class TestAgent:
    def __init__(self, invalid=False):
        self.active = 0
        self.peak = 0
        self.invalid = invalid

    async def run(self, prompt, **kwargs):
        self.active += 1
        self.peak = max(self.peak, self.active)
        try:
            await asyncio.sleep(0.001)
            plan = {'status': 'ready', 'actions': []} if self.invalid else solve(json.loads(prompt)['scene'])[0]
            if plan['actions']:
                scene = json.loads(prompt)['scene']
                plan['actions'] = [{**a, 'text': action_text(a, scene, Config().prompts)} for a in plan['actions']]
            return SimpleNamespace(output=Plan.model_validate(plan))
        finally:
            self.active -= 1


class PipelineTests(unittest.TestCase):
    def teacher_config(self, path):
        return Config(output_dir=Path(path), target_examples=30,
                      teacher=TeacherConfig(enabled=True, fraction=1, concurrency=2, retries=1,
                                            llm=OpenAiConfig(provider='openai', model='test')))

    def test_teacher_concurrency_and_validated_provenance(self):
        with tempfile.TemporaryDirectory() as tmp:
            config = self.teacher_config(tmp)
            agent = TestAgent()
            asyncio.run(generate(config, agent))
            report = export(config)
            self.assertEqual(agent.peak, 2)
            self.assertEqual(report['sources'], {'openai': 10})
            self.assertEqual(report['teacher_failures'], 0)
            self.assertEqual(validate(config)['examples_checked'], 30)

    def test_invalid_teacher_cannot_poison_corpus(self):
        with tempfile.TemporaryDirectory() as tmp:
            config = self.teacher_config(tmp)
            asyncio.run(generate(config, TestAgent(invalid=True)))
            report = export(config)
            self.assertEqual(report['teacher_failures'], 10)
            self.assertEqual(report['sources'], {'symbolic': 10})
            self.assertEqual(validate(config)['examples_checked'], 30)

    def test_restart_after_interrupted_generation_equals_clean_run(self):
        import assembly_dataset.pipeline as pipeline
        original = pipeline.make_episode
        def interrupted(index, config):
            if index >= 40:
                raise KeyboardInterrupt()
            return original(index, config)
        with tempfile.TemporaryDirectory() as first, tempfile.TemporaryDirectory() as second:
            config = Config(output_dir=Path(first), target_examples=150)
            with patch.object(pipeline, 'make_episode', side_effect=interrupted):
                with self.assertRaises(KeyboardInterrupt):
                    asyncio.run(generate(config))
            asyncio.run(generate(config))
            resumed = export(config)
            clean = Config(output_dir=Path(second), target_examples=150)
            asyncio.run(generate(clean))
            uninterrupted = export(clean)
            self.assertEqual(resumed['files'], uninterrupted['files'])
            self.assertEqual(validate(config)['examples_checked'], 150)

    def test_corrupted_shard_is_detected(self):
        with tempfile.TemporaryDirectory() as tmp:
            config = Config(output_dir=Path(tmp), target_examples=4)
            asyncio.run(generate(config))
            export(config)
            next(Path(tmp).glob('sft/*/*/*.jsonl')).write_text('broken\n')
            with self.assertRaisesRegex(ValueError, 'Повреждён'):
                validate(config)

    def test_all_steps_and_one_mode_obey_total(self):
        for modes in (['step'], ['branch'], ['full', 'step', 'branch']):
            with self.subTest(modes=modes), tempfile.TemporaryDirectory() as tmp:
                config = Config(output_dir=Path(tmp), target_examples=37, modes=modes, step_sampling='all')
                asyncio.run(generate(config))
                report = export(config)
                self.assertEqual(report['examples'], 37)
                self.assertEqual(validate(config)['examples_checked'], 37)

    def test_smaller_export_does_not_remove_canonical_episodes(self):
        with tempfile.TemporaryDirectory() as tmp:
            config = Config(output_dir=Path(tmp), target_examples=100)
            asyncio.run(generate(config))
            original = export(config)
            config.target_examples = 1
            self.assertEqual(export(config)['examples'], 1)
            config.target_examples = 100
            self.assertEqual(export(config)['files'], original['files'])


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