from __future__ import annotations

from typing import Any, Literal, Optional

from pydantic import BaseModel, ConfigDict, Field


class TrainingDatasetCreateRequest(BaseModel):
    model_config = ConfigDict(extra="ignore")

    dataset_type: Literal["sft", "preference", "eval"]
    version: str
    asset_ids: list[int] = Field(default_factory=list)
    pair_ids: list[int] = Field(default_factory=list)
    selection_rules: dict[str, Any] = Field(default_factory=dict)
    notes: Optional[str] = None


class TrainingJobCreateRequest(BaseModel):
    model_config = ConfigDict(extra="ignore")

    dataset_id: int
    job_type: Literal["sft", "preference"]
    base_model_family: str
    target_slot: Literal["low", "high"]
    provider_backend: str


class TrainingDeploymentDecisionRequest(BaseModel):
    model_config = ConfigDict(extra="ignore")

    decision: Literal["accepted", "rejected", "rolled_back"]
    notes: Optional[str] = None