"""Base configuration schemas for BYOM train, eval, export, and deploy.

All schemas are extensible via Pydantic model inheritance. BYOM projects
add task-specific fields by subclassing these base models.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any, Literal

from pydantic import BaseModel, Field


class ModelInfo(BaseModel):
    """Identity and capabilities of a model — shared across all pipelines.

    Used by train, eval, export, deploy, and predict. Each BYOM project
    may subclass to add task-specific fields (e.g. conf_threshold for detection).
    """

    task: Literal["classification", "detection", "segmentation", "video_action"]
    architecture: str = Field(..., description="e.g. resnet50, yolov10-s, r2plus1d_18")
    num_classes: int = Field(..., gt=0, description="Number of output classes")
    input_size: tuple[int, ...] = Field(
        ...,
        min_length=2,
        max_length=4,
        description="(H, W) for images; (C, T, H, W) for video",
    )
    pretrained: bool = True
    """Extensible: detection adds conf_threshold, video adds num_frames."""


class TrainConfig(BaseModel):
    """Training hyperparameters — shared across all BYOM models.

    Projects may subclass to add task-specific fields (mixup, anchor config, etc.).
    """

    # -- Core --
    dataset_path: Path = Field(..., description="Root path to the training dataset")
    epochs: int = Field(gt=0, le=10000, description="Maximum training epochs")
    batch_size: int = Field(gt=0, le=1024, description="Training batch size")
    device: Literal["cuda", "cpu", "mps"] = Field(default="cuda", description="Training device")
    num_workers: int = Field(default=4, ge=0, description="DataLoader worker processes")

    # -- Optimiser --
    optimizer: str = Field(
        default="AdamW",
        description="Optimizer class name (any torch.optim class, e.g. SGD, Adam, AdamW)",
    )
    learning_rate: float = Field(gt=0.0, le=1.0, description="Initial learning rate")
    weight_decay: float = Field(default=1e-4, ge=0.0, description="L2 regularisation weight")
    momentum: float = Field(default=0.9, ge=0.0, le=1.0, description="SGD / RMSprop momentum")
    optimizer_kwargs: dict[str, Any] = Field(
        default_factory=dict,
        description="Extra kwargs passed to the optimizer constructor (override defaults or add new params)",
    )

    # -- LR scheduler --
    scheduler: str | None = Field(
        default=None,
        description="Scheduler class name (any torch.optim.lr_scheduler class, e.g. StepLR, CosineAnnealingLR)",
    )
    lr_step_size: int = Field(default=30, gt=0, description="Epoch interval for StepLR")
    lr_gamma: float = Field(default=0.1, gt=0.0, le=1.0, description="Multiplicative factor for StepLR / ExponentialLR")
    lr_min: float = Field(default=0.0, ge=0.0, description="Minimum LR for CosineAnnealingLR")
    scheduler_kwargs: dict[str, Any] = Field(
        default_factory=dict,
        description="Extra kwargs passed to the scheduler constructor (override defaults or add new params)",
    )

    # -- GPU / precision --
    use_amp: bool = Field(
        default=False,
        description="Enable automatic mixed precision (AMP) with GradScaler for CUDA training",
    )
    clear_cache_per_epoch: bool = Field(
        default=False,
        description="Call torch.cuda.empty_cache() after every epoch to release unused GPU memory",
    )

    # -- Early stopping --
    early_stopping_patience: int | None = Field(
        default=None,
        ge=1,
        description="Stop if no improvement for N epochs",
    )
    early_stopping_min_delta: float = Field(
        default=0.001,
        ge=0.0,
        description="Minimum metric change to qualify as improvement",
    )
    """Extensible: detection adds anchor config, video adds temporal settings."""


class EvalConfig(BaseModel):
    """Evaluation settings — framework and dataset.

    Used by eval.py. Projects may add task-specific metric thresholds.
    """

    dataset_path: Path = Field(..., description="Path to evaluation dataset")
    framework: Literal["pytorch", "onnx", "openvino", "tensorrt"] = Field(..., description="Inference framework")
    batch_size: int = Field(default=32, gt=0, le=1024)
    device: Literal["cuda", "cpu"] = Field(default="cuda")
    """Extensible: detection adds IoU thresholds, classification adds top_k."""


class ExportConfig(BaseModel):
    """Export pipeline configuration.

    Formats to export, validation, and upload targets.
    """

    formats: list[Literal["onnx", "tensorrt", "openvino", "torchscript"]] = Field(
        ..., min_length=1, description="Target export formats"
    )
    validate_after_export: bool = Field(default=True, description="Run validators after each format")
    upload_target: str | None = Field(default=None, description="Optional upload URI")


class DeployConfig(BaseModel):
    """Deployment server configuration."""

    port: int = Field(default=8080, ge=1, le=65535)
    framework: Literal["pytorch", "onnx", "openvino", "tensorrt"] = Field(
        default="pytorch", description="Serving framework"
    )
    model_path: Path | None = Field(default=None, description="Path to model file")
    device: Literal["cuda", "cpu"] = Field(default="cuda")


class ExperimentConfig(BaseModel):
    """Unified experiment config — single YAML with all sections.

    One file per experiment with nested model, training, eval, export,
    and deploy sections. Entry points access only what they need
    (e.g. ``config.model``, ``config.training``).
    """

    model: ModelInfo
    training: TrainConfig
    eval: EvalConfig
    export: ExportConfig
    deploy: DeployConfig = Field(default_factory=DeployConfig)
