"""Editable prompts and validated payload templates."""
from pathlib import Path
from string import Template, Formatter

import yaml
from pydantic import model_validator
from llm.config import ConfigModel

DEFAULT_PROMPTS = Path(__file__).resolve().parents[1] / 'prompts.yaml'


class Prompts(ConfigModel):
    teacher_system: str
    teacher_user: str
    teacher_retry: str
    training_system: dict[str, str]
    structured_system: dict[str, str]
    training_user: str
    task_instructions: dict[str, str]
    object_names: dict[str, str]
    action_templates: dict[str, str]

    @model_validator(mode='after')
    def validate_texts(self):
        texts = [self.teacher_system, *self.training_system.values(), *self.structured_system.values(), *self.task_instructions.values()]
        if any(not text.strip() for text in texts):
            raise ValueError('Промпты не должны быть пустыми')
        if set(self.training_system) != {'full', 'step', 'branch'}:
            raise ValueError('training_system должен содержать full, step, branch')
        if set(self.structured_system) != {'full', 'step', 'branch'}:
            raise ValueError('structured_system должен содержать full, step, branch')
        for value in (self.teacher_user, self.teacher_retry, self.training_user):
            template = Template(value)
            if not template.is_valid() or set(template.get_identifiers()) != {'payload'}:
                raise ValueError('Шаблон должен содержать ${payload}; другие переменные не поддерживаются')
        for value in self.action_templates.values():
            if not value.strip():
                raise ValueError('Пустой шаблон действия')
            fields = {field for _, field, _, _ in Formatter().parse(value) if field is not None}
            if fields - {'object', 'target'}:
                raise ValueError('В action_templates допустимы только {object} и {target}')
            value.format(object='объект', target='место')
        return self


def load_prompts(path: Path = DEFAULT_PROMPTS) -> Prompts:
    return Prompts.model_validate(yaml.safe_load(path.read_text(encoding='utf-8')))


def render(template: str, payload: str) -> str:
    return Template(template).substitute(payload=payload)