# ta_ml/models/api.py
"""
Pydantic models for API request and response validation.

This module centralizes all data transfer objects (DTOs) used by the FastAPI
application, ensuring consistent and validated data structures across the API.
"""
from typing import List, Optional

from pydantic import BaseModel, Field


# --- Prediction Models ---
class PredictRequest(BaseModel):
    """Request model for the prediction endpoint."""
    text: str = Field(..., min_length=1,
                      description="The transaction text to classify.")


class PredictResponse(BaseModel):
    """Response model for a successful prediction."""
    predicted_label: str = Field(..., description="The most likely category.")
    confidence: float = Field(..., ge=0, le=1,
                              description="The model's confidence score.")


# --- Tokenization Models ---
class TokenizationResult(BaseModel):
    """Detailed tokenization result for a single word or text."""
    word: str
    tokens: List[str]
    token_ids: List[int]
    token_vocab_flags: List[bool]
    all_tokens_in_vocab: bool
    vocab_status: str


class TokenizationDatasetStats(BaseModel):
    """Statistics for a tokenized dataset."""
    total_entries_processed: int
    in_vocab_count: int
    oov_token_count: int
    total_token_count: int
    in_vocab_ratio: float
    oov_token_ratio: float


class TokenizationDatasetResponse(BaseModel):
    """Response model for the file tokenization endpoint."""
    stats: TokenizationDatasetStats
    results: List[TokenizationResult]


# --- Training Models ---
class TrainDirectRequest(BaseModel):
    """Request model for training directly from text/label lists."""
    texts: List[str] = Field(..., min_items=1)
    labels: List[str] = Field(..., min_items=1)


class TrainingResponse(BaseModel):
    """Standard response for initiating a training job."""
    message: str


class TrainingMetrics(BaseModel):
    """Detailed metrics from a completed training process."""
    accuracy: float
    precision: float
    recall: float
    f1_score: float
    details: str


class TrainingCompletionResponse(BaseModel):
    """Response model for a completed training job."""
    message: str
    metrics: TrainingMetrics


# --- Backup Models ---
class CacheFileListResponse(BaseModel):
    """Response model for listing files in the cache."""
    files: List[str]


class BackupUploadResponse(BaseModel):
    """Response model for a successful cache upload."""
    status: str
    repo: str


# --- Admin Models ---
class ReloadResponse(BaseModel):
    """Response model for a model reload request."""
    status: str
    detail: Optional[str] = None