"""Admin Training Schemas."""

from datetime import datetime
from typing import Any

from pydantic import BaseModel, Field

from .augmentation import AugmentationConfigSchema
from .enums import TrainingStatus, TrainingType


class TrainingConfig(BaseModel):
    """Training configuration."""

    model_name: str = Field(default="yolo26s.pt", description="Base model name (used if no base_model_version_id)")
    base_model_version_id: str | None = Field(
        default=None,
        description="Model version UUID to use as base for incremental training. If set, uses this model instead of model_name.",
    )
    epochs: int = Field(default=100, ge=1, le=1000, description="Training epochs")
    batch_size: int = Field(default=16, ge=1, le=128, description="Batch size")
    image_size: int = Field(default=640, ge=320, le=1280, description="Image size")
    learning_rate: float = Field(default=0.01, gt=0, le=1, description="Learning rate")
    device: str = Field(default="0", description="Device (0 for GPU, cpu for CPU)")
    project_name: str = Field(
        default="invoice_fields", description="Training project name"
    )

    # Data augmentation settings
    augmentation: AugmentationConfigSchema | None = Field(
        default=None,
        description="Augmentation configuration. If provided, augments dataset before training.",
    )
    augmentation_multiplier: int = Field(
        default=2,
        ge=1,
        le=10,
        description="Number of augmented copies per original image",
    )


class TrainingTaskCreate(BaseModel):
    """Request to create a training task."""

    name: str = Field(..., min_length=1, max_length=255, description="Task name")
    description: str | None = Field(None, max_length=1000, description="Description")
    task_type: TrainingType = Field(
        default=TrainingType.TRAIN, description="Task type"
    )
    config: TrainingConfig = Field(
        default_factory=TrainingConfig, description="Training configuration"
    )
    scheduled_at: datetime | None = Field(
        None, description="Scheduled execution time"
    )
    cron_expression: str | None = Field(
        None, max_length=50, description="Cron expression for recurring tasks"
    )


class TrainingTaskItem(BaseModel):
    """Single training task in list."""

    task_id: str = Field(..., description="Task UUID")
    name: str = Field(..., description="Task name")
    task_type: TrainingType = Field(..., description="Task type")
    status: TrainingStatus = Field(..., description="Task status")
    scheduled_at: datetime | None = Field(None, description="Scheduled time")
    is_recurring: bool = Field(default=False, description="Is recurring task")
    started_at: datetime | None = Field(None, description="Start time")
    completed_at: datetime | None = Field(None, description="Completion time")
    created_at: datetime = Field(..., description="Creation timestamp")


class TrainingTaskListResponse(BaseModel):
    """Response for training task list."""

    total: int = Field(..., ge=0, description="Total tasks")
    limit: int = Field(..., ge=1, description="Page size")
    offset: int = Field(..., ge=0, description="Current offset")
    tasks: list[TrainingTaskItem] = Field(default_factory=list, description="Task list")


class TrainingTaskDetailResponse(BaseModel):
    """Response for training task detail."""

    task_id: str = Field(..., description="Task UUID")
    name: str = Field(..., description="Task name")
    description: str | None = Field(None, description="Description")
    task_type: TrainingType = Field(..., description="Task type")
    status: TrainingStatus = Field(..., description="Task status")
    config: dict[str, Any] | None = Field(None, description="Training configuration")
    scheduled_at: datetime | None = Field(None, description="Scheduled time")
    cron_expression: str | None = Field(None, description="Cron expression")
    is_recurring: bool = Field(default=False, description="Is recurring task")
    started_at: datetime | None = Field(None, description="Start time")
    completed_at: datetime | None = Field(None, description="Completion time")
    error_message: str | None = Field(None, description="Error message")
    result_metrics: dict[str, Any] | None = Field(None, description="Result metrics")
    model_path: str | None = Field(None, description="Trained model path")
    created_at: datetime = Field(..., description="Creation timestamp")


class TrainingTaskResponse(BaseModel):
    """Response for training task operation."""

    task_id: str = Field(..., description="Task UUID")
    status: TrainingStatus = Field(..., description="Task status")
    message: str = Field(..., description="Status message")


class TrainingLogItem(BaseModel):
    """Single training log entry."""

    level: str = Field(..., description="Log level")
    message: str = Field(..., description="Log message")
    details: dict[str, Any] | None = Field(None, description="Additional details")
    created_at: datetime = Field(..., description="Timestamp")


class TrainingLogsResponse(BaseModel):
    """Response for training logs."""

    task_id: str = Field(..., description="Task UUID")
    logs: list[TrainingLogItem] = Field(default_factory=list, description="Log entries")


class ExportRequest(BaseModel):
    """Request to export annotations."""

    format: str = Field(
        default="yolo", description="Export format (yolo, coco, voc)"
    )
    include_images: bool = Field(
        default=True, description="Include images in export"
    )
    split_ratio: float = Field(
        default=0.8, ge=0.5, le=1.0, description="Train/val split ratio"
    )


class ExportResponse(BaseModel):
    """Response for export operation."""

    status: str = Field(..., description="Export status")
    export_path: str = Field(..., description="Path to exported dataset")
    total_images: int = Field(..., ge=0, description="Total images exported")
    total_annotations: int = Field(..., ge=0, description="Total annotations")
    train_count: int = Field(..., ge=0, description="Training set count")
    val_count: int = Field(..., ge=0, description="Validation set count")
    message: str = Field(..., description="Status message")


class TrainingDocumentItem(BaseModel):
    """Document item for training page."""

    document_id: str = Field(..., description="Document UUID")
    filename: str = Field(..., description="Filename")
    annotation_count: int = Field(..., ge=0, description="Total annotations")
    annotation_sources: dict[str, int] = Field(
        ..., description="Annotation counts by source (manual, auto)"
    )
    used_in_training: list[str] = Field(
        default_factory=list, description="List of training task IDs that used this document"
    )
    last_modified: datetime = Field(..., description="Last modification time")


class TrainingDocumentsResponse(BaseModel):
    """Response for GET /admin/training/documents."""

    total: int = Field(..., ge=0, description="Total document count")
    limit: int = Field(..., ge=1, le=100, description="Page size")
    offset: int = Field(..., ge=0, description="Pagination offset")
    documents: list[TrainingDocumentItem] = Field(
        default_factory=list, description="Documents available for training"
    )


class ModelMetrics(BaseModel):
    """Training model metrics."""

    mAP: float | None = Field(None, ge=0.0, le=1.0, description="Mean Average Precision")
    precision: float | None = Field(None, ge=0.0, le=1.0, description="Precision")
    recall: float | None = Field(None, ge=0.0, le=1.0, description="Recall")


class TrainingModelItem(BaseModel):
    """Trained model item for model list."""

    task_id: str = Field(..., description="Training task UUID")
    name: str = Field(..., description="Model name")
    status: TrainingStatus = Field(..., description="Training status")
    document_count: int = Field(..., ge=0, description="Documents used in training")
    created_at: datetime = Field(..., description="Creation timestamp")
    completed_at: datetime | None = Field(None, description="Completion timestamp")
    metrics: ModelMetrics = Field(..., description="Model metrics")
    model_path: str | None = Field(None, description="Path to model weights")
    download_url: str | None = Field(None, description="Download URL for model")


class TrainingModelsResponse(BaseModel):
    """Response for GET /admin/training/models."""

    total: int = Field(..., ge=0, description="Total model count")
    limit: int = Field(..., ge=1, le=100, description="Page size")
    offset: int = Field(..., ge=0, description="Pagination offset")
    models: list[TrainingModelItem] = Field(
        default_factory=list, description="Trained models"
    )


class TrainingHistoryItem(BaseModel):
    """Training history for a document."""

    task_id: str = Field(..., description="Training task UUID")
    name: str = Field(..., description="Training task name")
    trained_at: datetime = Field(..., description="Training timestamp")
    model_metrics: ModelMetrics | None = Field(None, description="Model metrics")
