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

from assembly_dataset.config import Config, load_config
from assembly_dataset.pipeline import generate
from assembly_dataset.world import make_episode, training_rows
from assembly_dataset.language import action_text


class PromptTests(unittest.TestCase):
    def test_custom_request_and_retry_reach_teacher(self):
        from assembly_dataset.teacher import enrich, Plan
        from llm.config import OpenAiConfig
        config = Config(blocked_fraction=0)
        config.teacher.fraction = 1.0
        config.teacher.llm = OpenAiConfig(provider='openai', model='test')
        config.prompts.teacher_user = 'REQUEST ${payload}'
        config.prompts.teacher_retry = 'REPAIR ${payload}'
        episode = make_episode(0, config)
        agent = SimpleNamespace(run=AsyncMock(side_effect=[
            SimpleNamespace(output=Plan(status='ready', actions=[])),
            SimpleNamespace(output=Plan.model_validate({**episode['plan'], 'actions': [{**a, 'text': action_text(a, episode['scene'], config.prompts)} for a in episode['plan']['actions']]})),
        ]))
        result, error = asyncio.run(enrich(episode, config.teacher, agent, config.prompts))
        self.assertIsNone(error)
        first, second = [call.args[0] for call in agent.run.call_args_list]
        self.assertTrue(first.startswith('REQUEST {'))
        self.assertTrue(second.startswith('REPAIR {'))
        self.assertIn('validation_error', json.loads(second.removeprefix('REPAIR ')))
        self.assertEqual(result['source']['provider'], 'openai')

    def test_external_prompts_relative_paths_and_export(self):
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            prompts = yaml.safe_load(Path('prompts.yaml').read_text())
            prompts['training_system']['full'] = 'Custom training rules'
            prompts['training_user'] = 'Input: ${payload}'
            prompts['task_instructions']['mechanical'] = 'Custom assembly task'
            (root / 'texts.yaml').write_text(yaml.safe_dump(prompts))
            (root / 'config.yaml').write_text('prompts_file: texts.yaml\noutput_dir: data\ndomains: [mechanical]\n')
            config = load_config(root / 'config.yaml')
            episode = make_episode(0, config)
            record = next(training_rows(episode, config))
            self.assertEqual(episode['instruction'], 'Custom assembly task')
            self.assertEqual(record['messages'][0]['content'], 'Custom training rules')
            self.assertTrue(record['messages'][1]['content'].startswith('Input: {'))
            self.assertEqual(config.output_dir, root / 'data')

    def test_prompt_changes_prevent_mixed_resume(self):
        with tempfile.TemporaryDirectory() as tmp:
            config = Config(output_dir=Path(tmp), target_examples=1)
            asyncio.run(generate(config))
            config.prompts.teacher_system = 'Different rules'
            with self.assertRaises(ValueError):
                asyncio.run(generate(config))

    def test_invalid_templates_are_rejected(self):
        from assembly_dataset.prompts import Prompts
        data = Config().prompts.model_dump()
        for template in ('No scene', '${unknown}', '${payload.bad}'):
            with self.subTest(template=template), self.assertRaises(ValueError):
                Prompts.model_validate({**data, 'teacher_user': template})

    def test_teacher_system_is_passed_to_agent(self):
        from assembly_dataset.teacher import create_agent
        from llm.config import OpenAiConfig
        config = Config()
        config.teacher.llm = OpenAiConfig(provider='openai', model='test')
        config.prompts.teacher_system = 'Custom teacher instructions'
        with patch('llm.factory.LLMModelProvider.create_from_config', new_callable=AsyncMock), patch('pydantic_ai.Agent') as agent:
            asyncio.run(create_agent(config.teacher, config.prompts))
            self.assertEqual(agent.call_args.kwargs['system_prompt'], 'Custom teacher instructions')