from pathlib import Path
from typing import Literal

import yaml
from pydantic import Field, model_validator

from decision_lab.config import StrictModel


class PilotConfig(StrictModel):
    model_id: str
    revision: str = Field(pattern=r"^[0-9a-f]{40}$")
    seed: int
    max_length: int = Field(gt=0)
    class_tokens: list[str]
    class_labels: list[str]
    answer_prefix: str
    microbatch: int = Field(gt=0)
    steps: int = Field(gt=0)
    warmup_steps: int = Field(ge=0)
    examples: int = Field(gt=0)
    backbone_lr: float = Field(gt=0)
    head_lr: float = Field(gt=0)
    weight_decay: float = Field(ge=0)
    gradient_clip: float = Field(gt=0)
    cpu_threads: int = Field(gt=0)
    gpu: Literal["L4"]
    cpu_cores: int = Field(gt=0)
    memory_mib: int = Field(gt=0)
    timeout_seconds: int = Field(gt=0)
    startup_timeout_seconds: int = Field(gt=0)
    scaledown_seconds: int = Field(ge=0)
    gpu_hourly_usd: float = Field(gt=0)
    cpu_core_hourly_usd: float = Field(gt=0)
    memory_gib_hourly_usd: float = Field(gt=0)
    max_pilot_usd: float = Field(gt=0)
    required_remaining_usd: float = Field(gt=0)
    volume_name: str
    app_name: str

    @model_validator(mode="after")
    def validate_classes(self):
        if self.class_labels != ["TRUE", "FALSE", "UNKNOWN"]:
            raise ValueError("The shared class order must be TRUE, FALSE, UNKNOWN")
        if len(set(self.class_tokens)) != len(self.class_labels):
            raise ValueError("Each class needs a distinct answer token")
        if self.warmup_steps >= self.steps:
            raise ValueError("Warmup steps must be less than total steps")
        return self

    @property
    def hourly_usd(self) -> float:
        return (
            self.gpu_hourly_usd
            + self.cpu_cores * self.cpu_core_hourly_usd
            + self.memory_mib / 1024 * self.memory_gib_hourly_usd
        )


class ScratchConfig(StrictModel):
    vocab_size: int = Field(gt=0)
    hidden_size: int = Field(gt=0)
    heads: int = Field(gt=0)
    layers: int = Field(gt=0)
    intermediate_size: int = Field(gt=0)
    max_length: int = Field(gt=0)
    dropout: float = Field(ge=0, lt=1)
    classes: Literal[3]
    pad_id: int = Field(ge=0)
    initializer_std: float = Field(gt=0)

    @model_validator(mode="after")
    def validate_shape(self):
        if self.hidden_size % self.heads:
            raise ValueError("Hidden dimension must be divisible by attention heads")
        if self.pad_id >= self.vocab_size:
            raise ValueError("Padding ID must be within the vocabulary")
        return self


def load_pilot_config(path: Path) -> PilotConfig:
    return PilotConfig.model_validate(yaml.safe_load(path.read_text()))


def load_scratch_config(path: Path) -> ScratchConfig:
    return ScratchConfig.model_validate(yaml.safe_load(path.read_text()))