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", "eval"]
    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


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

    base_model_family: str = Field(default="routing_default", min_length=1, max_length=128)
    provider_backend: str = Field(default="control-plane-training", min_length=1, max_length=128)
    target_slots: list[Literal["low", "high"]] = Field(default_factory=lambda: ["low", "high"])
    notes: Optional[str] = None


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

    role: str
    dataset_id: int
    dataset_type: str
    version: str
    sample_count: int = 0
    suite_kind: Optional[str] = None


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

    role: str
    job_id: int
    job_type: str
    dataset_id: int
    target_slot: str
    status: str
    artifact_id: Optional[int] = None


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

    ready_for_acceptance: bool = False
    threshold_score: float = 90.0
    frozen_eval_score: Optional[float] = None
    severe_blocker_count: int = 0
    blocking_reason: str


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

    version_prefix: str
    datasets: list[MountedTrainingWaveDatasetItem] = Field(default_factory=list)
    jobs: list[MountedTrainingWaveJobItem] = Field(default_factory=list)
    acceptance_gate: MountedTrainingAcceptanceGate