import uuid
from datetime import datetime

from pydantic import BaseModel, ConfigDict

from app.models.stubs import FineTuningJobStatus


class FineTuningJobCreate(BaseModel):
    dataset_id: uuid.UUID
    base_model: str
    name: str
    ollama_tag: str | None = None
    system_prompt_id: uuid.UUID | None = None
    iters: int = 100
    learning_rate: float = 1e-5
    batch_size: int = 2
    num_layers: int = 16  # MLX backend only
    lora_r: int = 16  # CUDA backend only
    lora_alpha: int = 16  # CUDA backend only
    lora_dropout: float = 0.0  # CUDA backend only


class FineTuningJobOut(BaseModel):
    model_config = ConfigDict(from_attributes=True)

    id: uuid.UUID
    project_id: uuid.UUID
    dataset_id: uuid.UUID
    model_version_id: uuid.UUID | None
    base_model: str
    config: dict
    status: FineTuningJobStatus
    progress: dict
    created_at: datetime