"""Typed configuration for all training stages, loaded from YAML."""

from dataclasses import dataclass, field, fields
from pathlib import Path

import yaml


@dataclass
class ModelConfig:
    """Which pretrained model to load."""

    name_or_path: str
    tokenizer: str | None = None  # defaults to name_or_path if None
    trust_remote_code: bool = False


@dataclass
class DAPTDataConfig:
    """Data paths and preprocessing for DAPT/TAPT."""

    corpus_path: str  # directory of JSONL shards or single JSONL file
    text_field: str = "text"
    max_seq_length: int = 2048
    max_tokens: int | None = None  # subsample corpus to ~N tokens (None = use all)
    validation_split: float = 0.02


@dataclass
class TrainingConfig:
    """HuggingFace Trainer arguments."""

    output_dir: str
    learning_rate: float = 5e-5
    mlm_probability: float = 0.30
    whole_word_mask: bool = False
    num_train_epochs: int = 1
    per_device_train_batch_size: int = 4
    gradient_accumulation_steps: int = 8
    warmup_ratio: float = 0.05
    weight_decay: float = 0.01
    bf16: bool = True
    gradient_checkpointing: bool = True
    logging_steps: int = 50
    save_strategy: str = "steps"  # "steps" or "epoch"
    save_steps: int = 1000  # ignored when save_strategy="epoch"
    eval_strategy: str = "steps"  # "steps" or "epoch"
    eval_steps: int = 1000  # ignored when eval_strategy="epoch"
    save_total_limit: int = 3
    dataloader_num_workers: int = 4
    seed: int = 42
    resume_from_checkpoint: str | None = None


@dataclass
class DAPTConfig:
    """Full configuration for a DAPT or TAPT run."""

    model: ModelConfig
    data: DAPTDataConfig
    training: TrainingConfig
    stage: str = "dapt"  # "dapt" or "tapt" — informational label

    @classmethod
    def from_yaml(cls, path: str | Path) -> "DAPTConfig":
        with open(path) as f:
            raw = yaml.safe_load(f)
        return cls(
            model=ModelConfig(**raw["model"]),
            data=DAPTDataConfig(**raw["data"]),
            training=TrainingConfig(**raw["training"]),
            stage=raw.get("stage", "dapt"),
        )

    def apply_overrides(
        self,
        *,
        model_path: str | None = None,
        data_path: str | None = None,
        output_dir: str | None = None,
        stage: str | None = None,
    ) -> None:
        """Apply CLI overrides on top of YAML config."""
        if model_path is not None:
            self.model.name_or_path = model_path
        if data_path is not None:
            self.data.corpus_path = data_path
        if output_dir is not None:
            self.training.output_dir = output_dir
        if stage is not None:
            self.stage = stage
