"""Optional teacher plans through the project's existing model factory."""
import asyncio
import re
from typing import Literal

from pydantic import Field, field_validator
from pydantic_ai.exceptions import UnexpectedModelBehavior
from llm.config import ConfigModel
from .prompts import render
from .language import validate_texts, action_text
from .world import canonical, validate_plan, branch_plan, validate_branch


class Action(ConfigModel):
    skill: str
    object: str
    target: str
    text: str = Field(description='Одно действие на русском языке с названием объекта из сцены; без кода')

    @field_validator('text')
    @classmethod
    def russian_description(cls, value):
        value = value.strip()
        if not re.search('[А-Яа-яЁё]', value) or re.search(r'\b(?:part_|slot_)\d+|\b(?:robot|camera)\.', value):
            raise ValueError('Ожидается русское описание действия, без технических идентификаторов')
        return value


class Plan(ConfigModel):
    status: Literal['ready', 'blocked']
    actions: list[Action] = Field(default_factory=list)


async def create_agent(config, prompts):
    from pydantic_ai import Agent
    from llm.factory import LLMModelProvider
    model = await LLMModelProvider.create_from_config(config.llm)
    return Agent(model=model, output_type=Plan, retries=config.retries,
                 system_prompt=prompts.teacher_system)


async def enrich(episode, config, agent, prompts):
    if int(episode['signature'][16:32], 16) / 2**64 >= config.fraction:
        return episode, None
    teacher_scene = {**episode['scene'], 'operations': [
        {**operation, 'reference_text': action_text(operation['action'], episode['scene'], prompts)}
        for operation in episode['scene']['operations']]}
    prompt = render(prompts.teacher_user, canonical({'instruction': episode['instruction'], 'scene': teacher_scene}))
    error = None
    for attempt in range(config.retries + 1):
        try:
            async with asyncio.timeout(config.timeout_seconds):
                result = await agent.run(prompt, model_settings=config.generation.model_dump(exclude_none=True))
            raw = result.output.model_dump()
            texts = [action['text'] for action in raw['actions']]
            plan = {'status': raw['status'], 'actions': [{k: v for k, v in a.items() if k != 'text'} for a in raw['actions']]}
            validate_texts({**episode, 'plan': plan}, texts, prompts)
            trace = validate_plan(episode, plan)
            branch = branch_plan(plan)
            validate_branch(episode, branch)
            episode.update(plan=plan, trace=trace, branch=branch, step_texts=texts, text_source='teacher',
                           source={'provider': config.llm.provider, 'model': config.llm.model})
            return episode, None
        except Exception as exc:
            # API exceptions may contain request headers. Persist the class only.
            error = {'scenario_id': episode['scenario_id'], 'error_type': type(exc).__name__, 'attempts': attempt + 1}
            if not isinstance(exc, (ValueError, UnexpectedModelBehavior)) and attempt == config.retries:
                raise RuntimeError(f'Запрос LLM не выполнен: {type(exc).__name__}. Генерация остановлена.') from None
            if isinstance(exc, ValueError):
                prompt = render(prompts.teacher_retry, canonical({'instruction': episode['instruction'], 'scene': teacher_scene,
                                    'validation_error': str(exc)[:500]}))
    # A failed teacher is never silently labelled as a successful teacher plan.
    episode['teacher_error'] = error
    return episode, error