"""Pydantic models for the optimizer."""

from typing import Literal

from pydantic import BaseModel, Field

from bilgeai.core.constants import SCHEMA_VERSION


class Explanation(BaseModel):
    """First-class data type: every decision produces an Explanation."""

    decision: str  # "batch_size = 4"
    reason: str  # "16GB VRAM, seq_len=512 average"
    confidence: float = Field(ge=0.0, le=1.0)  # 0.0 - 1.0
    source: Literal["rule_engine", "local_llm", "remote_api", "llm", "user"]
    alternatives: list[str] = Field(default_factory=list)


class HardwareConstraints(BaseModel):
    """Safety bounds computed by guardrail rules. Passed to LLM as context."""

    max_batch_size: int
    required_quantization: Literal["none", "4bit", "8bit"] | None = None
    max_seq_length_safe: int = 2048
    distributed_strategy: str | None = None
    num_gpus: int = 1
    gradient_checkpointing_required: bool = False
    explanations: list[Explanation] = Field(default_factory=list)


class TrainingConfig(BaseModel):
    """Complete training configuration with explanations."""

    schema_version: str = SCHEMA_VERSION
    model_name: str = "meta-llama/Llama-3-8B"
    quantization: Literal["none", "4bit", "8bit"] | None = None
    batch_size: int = Field(default=4, gt=0)
    gradient_accumulation_steps: int = Field(default=1, gt=0)
    learning_rate: float = Field(default=2e-4, gt=0)
    lr_scheduler: str = "cosine"
    warmup_ratio: float = Field(default=0.03, ge=0, le=1)
    num_epochs: int = Field(default=3, gt=0)
    max_seq_length: int = Field(default=512, gt=0)
    lora_enabled: bool = True
    lora_rank: int | None = 16
    lora_alpha: int | None = 32
    lora_target_modules: list[str] | None = None
    use_gradient_checkpointing: bool = False
    optimizer_name: str = "adamw_torch"
    fp16: bool = False
    bf16: bool = False

    # Distributed training
    num_gpus: int = 1
    num_nodes: int = 1
    distributed_strategy: Literal["none", "ddp", "fsdp", "deepspeed"] | None = None

    # DeepSpeed config
    deepspeed_stage: int | None = None
    deepspeed_offload_optimizer: bool = False
    deepspeed_offload_param: bool = False

    # Training strategy
    save_strategy: Literal["no", "steps", "epoch"] = "epoch"
    save_steps: int | None = None
    save_total_limit: int = 2
    logging_steps: int = 10
    packing: bool = False

    # Launch command (computed, not from LLM)
    torchrun_command: str | None = None

    # Reproducibility
    seed: int = 42

    # Experiment tracking
    report_to: Literal["none", "wandb", "mlflow", "tensorboard"] | None = None
    wandb_project: str | None = None

    explanations: list[Explanation] = Field(default_factory=list)

    def add_explanation(
        self,
        decision: str,
        reason: str,
        confidence: float = 0.9,
        source: str = "rule_engine",
        alternatives: list[str] | None = None,
    ) -> None:
        """Add an explanation for a decision."""
        self.explanations.append(
            Explanation(
                decision=decision,
                reason=reason,
                confidence=confidence,
                source=source,
                alternatives=alternatives or [],
            )
        )
