"""Typed dataclass containers mirroring the YAML configuration structure."""

from __future__ import annotations

from dataclasses import dataclass

__all__ = [
    "ModelCfg",
    "DatasetCfg",
    "OptimCfg",
    "SchedCfg",
    "TokenizerCfg",
    "TrainerCfg",
    "RuntimeCfg",
    "GenerationCfg",
    "TrainAppCfg",
    "InferAppCfg",
]


@dataclass
class ModelCfg:
    name: str
    best_checkpoint_path: str
    latest_checkpoint_path: str
    tokenizer: str
    d_model: int
    num_layers: int
    num_heads: int
    d_ff: int
    dropout_rate: float
    max_seq_len: int
    vocab_size: int
    pad_id: int
    bos_id: int
    eos_id: int
    layer_norm_style: str | None = None


@dataclass
class DatasetCfg:
    name: str
    path: str


@dataclass
class OptimCfg:
    name: str
    lr: float
    betas: tuple[float, float]
    weight_decay: float
    eps: float


@dataclass
class SchedCfg:
    name: str
    warmup_ratio: float
    min_lr_ratio: float
    max_lr_ratio: float
    hold_min: bool


@dataclass
class TokenizerCfg:
    name: str
    path: str
    corpus: str
    dataset: str
    vocab_size: int
    max_seq_len: int
    pad_token: str
    bos_token: str
    eos_token: str
    unk_token: str
    special_tokens: list[str]
    pad_id: int
    bos_id: int
    eos_id: int
    unk_id: int


@dataclass
class TrainerCfg:
    epochs: int
    batch_size: int
    shuffle: bool
    precision: str
    gradient_accumulation: int
    clip_grad_norm: float
    eval_interval: int
    save_interval: int
    debug_interval: int
    mixed_precision: bool
    gradient_checkpointing: bool
    label_smoothing: float
    export_best_to_model_path: bool
    resume: str


@dataclass
class RuntimeCfg:
    seed: int
    device: str
    output_dir: str
    data_dir: str
    cache_dir: str
    tokenizer_dir: str
    checkpoint_path: str


@dataclass
class GenerationCfg:
    max_new_tokens: int
    temperature: float
    top_k: int | None  # None disables top-k
    top_p: float | None
    do_sample: bool
    presence_penalty: float
    frequency_penalty: float
    no_repeat_ngram: int | None
    min_steps_before_eos: int


@dataclass
class TrainAppCfg:
    model: ModelCfg
    dataset: DatasetCfg
    optim: OptimCfg
    sched: SchedCfg
    trainer: TrainerCfg
    runtime: RuntimeCfg
    tokenizer: TokenizerCfg


@dataclass
class InferAppCfg:
    model: ModelCfg
    runtime: RuntimeCfg
    tokenizer: TokenizerCfg
    generation: GenerationCfg
    input_text: str