"""
Pydantic models for API requests and responses.
"""

from typing import Optional, Dict, Any, Literal
from pydantic import BaseModel, Field, validator
from datetime import datetime
import config


class GenerationParams(BaseModel):
    """Optional generation parameters."""
    width: int = Field(default=config.DEFAULT_WIDTH, ge=512, le=2048)
    height: int = Field(default=config.DEFAULT_HEIGHT, ge=512, le=2048)
    num_inference_steps: int = Field(default=config.DEFAULT_STEPS, ge=config.MIN_STEPS, le=config.MAX_STEPS)
    guidance_scale: float = Field(default=config.DEFAULT_GUIDANCE_SCALE, ge=config.MIN_GUIDANCE, le=config.MAX_GUIDANCE)
    seed: Optional[int] = Field(default=None, description="Random seed for reproducibility")
    face_image: Optional[str] = Field(default=None, description="Face reference image name (from faces directory) or 'vixy' for default")
    face_strength: float = Field(default=config.DEFAULT_FACE_STRENGTH, ge=0.0, le=1.0, description="Face conditioning strength (0.0-1.0)")

    @validator('width', 'height')
    def must_be_multiple_of_8(cls, v):
        if v % 8 != 0:
            raise ValueError('Width and height must be multiples of 8')
        return v


class GenerateRequest(BaseModel):
    """Request to generate an image."""
    prompt: str = Field(..., min_length=1, max_length=2000, description="Text prompt for image generation")
    negative_prompt: Optional[str] = Field(default=None, max_length=2000, description="Negative prompt to avoid certain features")
    params: Optional[GenerationParams] = Field(default_factory=GenerationParams)


class JobResponse(BaseModel):
    """Response when submitting a generation job."""
    job_id: str = Field(..., description="Unique job identifier")
    status: Literal["queued", "processing", "completed", "failed"] = Field(..., description="Current job status")
    created_at: datetime = Field(..., description="Job creation timestamp")
    message: Optional[str] = Field(default=None, description="Optional message")


class JobStatus(BaseModel):
    """Detailed job status information."""
    job_id: str
    status: Literal["queued", "processing", "completed", "failed"]
    progress: int = Field(..., ge=0, le=100, description="Progress percentage (0-100)")
    created_at: datetime
    started_at: Optional[datetime] = None
    completed_at: Optional[datetime] = None
    error: Optional[str] = None
    prompt: str


class HealthResponse(BaseModel):
    """Health check response."""
    model_config = {"protected_namespaces": ()}  # Allow "model_" prefix

    status: Literal["healthy", "unhealthy"]
    version: str
    model_loaded: bool
    queue_size: int
    active_jobs: int
    uptime_seconds: float


class ModelsResponse(BaseModel):
    """Available models information."""
    base_model: str
    refiner_model: Optional[str] = None
    refiner_enabled: bool
    device: str
    fp16_enabled: bool