from pathlib import Path
from typing import Annotated, Literal

import yaml
from pydantic import BaseModel, ConfigDict, Field


class ConfigModel(BaseModel):
    model_config = ConfigDict(extra="forbid")

# LLM
class CodexConfig(ConfigModel):
    provider: Literal["codex"]
    model: str = "gpt-5.6-luna"

    credential_path: Path | None = None


class OpenAiConfig(ConfigModel):
    provider: Literal["openai"]
    model: str

    api_key_env: str = "API_KEY"
    base_url_env: str = "BASE_URL"


LLMConfig = Annotated[
    CodexConfig | OpenAiConfig,
    Field(discriminator="provider")
]

# Generation
class GenerationConfig(ConfigModel):
    temperature: float = Field(default=.7,
                               ge=.0,
                               le=2.0)
    max_tokens: int = Field(default=4096,
                            gt=0)
    

# Runtime
class RuntimeConfig(ConfigModel):
    concurrency: int = Field(ge=1, default=1)

# Prompts
class PromptsConfig(ConfigModel):
    instructions: str
    system: str
    user: str

# Application
class AppConfig(ConfigModel):
    llm: LLMConfig
    generation: GenerationConfig = GenerationConfig()
    runtime: RuntimeConfig = RuntimeConfig()
    prompts: PromptsConfig


# Loader
def load_config(path: str | Path = "config.yaml") -> AppConfig:
    path = Path(path)

    with path.open("r", encoding="utf-8") as file:
        data = yaml.safe_load(file)

    return AppConfig.model_validate(data)