from pydantic import BaseModel, Field
from typing import List

class PredictionEntry(BaseModel):
    model: str = Field(..., description="Name of the model used for prediction")
    error: str | None = Field(None, description="Error message if prediction failed")
    predicted_class: int = Field(..., description="Predicted class index")
    confidence: float = Field(..., ge=-100.0, le=1.0, description="Prediction confidence (0-1)")

class ImagePredictionResponse(BaseModel):
    predictions: List[PredictionEntry] = Field(..., description="List of predictions with model, class, and confidence")

class TextPredictionRequest(BaseModel):
    text: str = Field(..., example="This is a sample text to classify")


class TextPredictionResponse(BaseModel):
    predicted_class: str = Field(..., description="Predicted label: 'AI' or 'Human'")
    confidence_ai: float = Field(..., ge=-100.0, le=1.0, description="Confidence score for AI (0-1)")
    confidence_human: float = Field(..., ge=-100.0, le=1.0, description="Confidence score for Human (0-1)")