from pathlib import Path
from typing import Literal

import yaml
from pydantic import Field, model_validator

from llm.config import ConfigModel, LLMConfig, GenerationConfig, CodexConfig

from .prompts import Prompts, load_prompts

DOMAINS = (
    'manipulation', 'mechanical', 'fastening', 'press_fit', 'electrical',
    'electronics', 'adhesive', 'welding', 'collaborative', 'inspection', 'operator_assistance',
)


class TeacherConfig(ConfigModel):
    enabled: bool = False
    fraction: float = Field(default=0.15, ge=0, le=1)
    llm: LLMConfig | None = None
    generation: GenerationConfig = Field(default_factory=GenerationConfig)
    concurrency: int = Field(default=4, ge=1, le=128)
    timeout_seconds: float = Field(default=180, gt=0)
    retries: int = Field(default=2, ge=0, le=10)

    @model_validator(mode='after')
    def require_model(self):
        if self.enabled and self.llm is None:
            raise ValueError('teacher.llm обязателен при teacher.enabled: true')
        return self


class Config(ConfigModel):
    # Count means exported SFT rows, not episodes or API calls.
    target_examples: int = Field(default=500_000, ge=0, strict=True)
    seed: int = 20261021
    output_dir: Path = Path('dataset/assembly')
    modes: list[Literal['full', 'step', 'branch']] = Field(default_factory=lambda: ['full', 'step', 'branch'], min_length=1)
    step_sampling: Literal['one', 'all'] = 'one'
    shard_size: int = Field(default=1000, ge=1)
    split: tuple[float, float, float] = (0.9, 0.05, 0.05)
    domains: list[str] = Field(default_factory=lambda: list(DOMAINS), min_length=1)
    min_components: int = Field(default=2, ge=1, le=12)
    max_components: int = Field(default=6, ge=1, le=12)
    blocked_fraction: float = Field(default=0.1, ge=0, le=1)
    prompts_file: Path = Path("prompts.yaml")
    prompts: Prompts = Field(default_factory=load_prompts)
    teacher: TeacherConfig = Field(default_factory=TeacherConfig)

    @model_validator(mode='after')
    def check(self):
        if len(set(self.modes)) != len(self.modes):
            raise ValueError('modes не должны повторяться')
        if len(set(self.domains)) != len(self.domains) or set(self.domains) - set(DOMAINS):
            raise ValueError(f'Допустимые domains: {DOMAINS}')
        if set(self.domains) - set(self.prompts.task_instructions):
            raise ValueError("В task_instructions отсутствуют инструкции выбранных domains")
        if self.min_components > self.max_components:
            raise ValueError('min_components > max_components')
        if any(v < 0 or v > 1 for v in self.split) or abs(sum(self.split) - 1) > 1e-9:
            raise ValueError('split должен содержать три доли с суммой 1')
        return self


def load_config(path: Path) -> Config:
    path = path.resolve()
    data = yaml.safe_load(path.read_text(encoding='utf-8')) or {}
    prompts_path = (path.parent / data.get('prompts_file', 'prompts.yaml')).resolve()
    data['prompts'] = load_prompts(prompts_path)
    data['prompts_file'] = prompts_path
    config = Config.model_validate(data)
    config.output_dir = (path.parent / config.output_dir).resolve()
    llm = config.teacher.llm
    if isinstance(llm, CodexConfig) and llm.credential_path:
        llm.credential_path = (path.parent / llm.credential_path).resolve()
    return config