import asyncio
import re
from types import SimpleNamespace
import unittest

from assembly_dataset.config import Config
from assembly_dataset.world import make_episode, training_rows, validate_plan


class RussianPlanTests(unittest.TestCase):
    def test_training_answers_are_russian_text_for_all_modes(self):
        config = Config(blocked_fraction=0)
        episode = make_episode(0, config)
        for row in training_rows(episode, config):
            text = row['messages'][-1]['content']
            self.assertRegex(text, '[А-Яа-я]')
            self.assertNotRegex(text, r'part_\d|slot_\d|"skill"|on_success')
            self.assertFalse(text.startswith('{'))
            if row['mode'] == 'full':
                self.assertTrue(text.startswith('1. '))
                self.assertIn('Захватить', text)

    def test_operator_assistance_holds_until_completion_then_returns_to_storage(self):
        config = Config(domains=['operator_assistance'], min_components=1, max_components=1, blocked_fraction=0)
        episode = make_episode(0, config)
        skills = [a['skill'] for a in episode['plan']['actions']]
        self.assertEqual(skills, ['detect', 'approach', 'pick', 'move_to_operator', 'hold',
                                  'wait_assembly', 'move_to_storage', 'release', 'return_home'])
        validate_plan(episode, episode['plan'])
        text = next(training_rows(episode, config))['messages'][-1]['content']
        self.assertIn('распределительная коробка', text)
        self.assertIn('завершения сборки', text)
        self.assertIn('хранения', text)
        self.assertIn('исходное положение', text)

    def test_teacher_wording_is_exported_without_replacing_it(self):
        from assembly_dataset.teacher import Plan, enrich
        from assembly_dataset.language import action_text
        from llm.config import OpenAiConfig
        config = Config(blocked_fraction=0)
        config.teacher.fraction = 1
        config.teacher.llm = OpenAiConfig(provider='openai', model='test')
        episode = make_episode(0, config)
        actions = [{**a, 'text': action_text(a, episode['scene'], config.prompts)} for a in episode['plan']['actions']]
        actions[0]['text'] = 'Сначала ' + actions[0]['text'].lower()
        output = Plan(status='ready', actions=actions)
        class Agent:
            async def run(self, *args, **kwargs):
                return SimpleNamespace(output=output)
        enriched, error = asyncio.run(enrich(episode, config.teacher, Agent(), config.prompts))
        self.assertIsNone(error)
        text = next(training_rows(enriched, config))['messages'][-1]['content']
        self.assertIn(actions[0]['text'], text)

    def test_teacher_cannot_supply_only_english_skill_text(self):
        from assembly_dataset.teacher import Plan
        with self.assertRaises(ValueError):
            Plan(status='ready', actions=[{'skill':'pick','object':'part_0','target':'slot_0','text':'pick: part_0 -> slot_0'}])

    def test_russian_text_cannot_change_the_validated_action(self):
        from assembly_dataset.language import action_text, validate_texts
        config = Config(domains=['operator_assistance'], min_components=1, max_components=1, blocked_fraction=0)
        episode = make_episode(0, config)
        texts = [action_text(a, episode['scene'], config.prompts) for a in episode['plan']['actions']]
        for bad in ('Бросить объект «распределительная коробка №1» на пол.',
                    texts[0] + ' Затем отпустить объект.'):
            with self.subTest(text=bad), self.assertRaises(ValueError):
                validate_texts(episode, [bad, *texts[1:]], config.prompts)

    def test_approaching_another_object_cannot_interrupt_holding(self):
        config = Config(domains=['operator_assistance'], min_components=2, max_components=2, blocked_fraction=0)
        episode = make_episode(0, config)
        actions = episode['plan']['actions']
        # Execute the second object's approach while the first is being held.
        reordered = actions[:5] + actions[9:11] + actions[5:9] + actions[11:]
        with self.assertRaises(ValueError):
            validate_plan(episode, {'status': 'ready', 'actions': reordered})