from pydantic import BaseModel, Field
from typing import List, Optional, Literal
from enum import Enum


class ImageSize(str, Enum):
    """Supported image sizes (OpenAI compatible)"""
    SMALL = "256x256"
    MEDIUM = "512x512" 
    LARGE = "1024x1024"
    WIDE = "1792x1024"
    TALL = "1024x1792"


class ImageQuality(str, Enum):
    """Image quality options"""
    STANDARD = "standard"
    HD = "hd"


class ImageStyle(str, Enum):
    """Image style options"""
    VIVID = "vivid"
    NATURAL = "natural"


class ResponseFormat(str, Enum):
    """Response format options"""
    URL = "url"
    B64_JSON = "b64_json"


class ImageGenerationRequest(BaseModel):
    """OpenAI compatible image generation request"""
    prompt: str = Field(..., description="A text description of the desired image(s)")
    model: str = Field(default="dall-e-3", description="The model to use for image generation")
    n: int = Field(default=1, ge=1, le=10, description="Number of images to generate")
    quality: ImageQuality = Field(default=ImageQuality.STANDARD, description="Quality of the image")
    response_format: ResponseFormat = Field(default=ResponseFormat.URL, description="Response format")
    size: ImageSize = Field(default=ImageSize.LARGE, description="Size of the generated images")
    style: ImageStyle = Field(default=ImageStyle.VIVID, description="Style of the generated images")
    user: Optional[str] = Field(default=None, description="A unique identifier representing your end-user")


class ImageData(BaseModel):
    """Individual image data in response"""
    url: Optional[str] = Field(default=None, description="URL of the generated image")
    b64_json: Optional[str] = Field(default=None, description="Base64 encoded image data")
    revised_prompt: Optional[str] = Field(default=None, description="The revised prompt used for generation")


class ImageGenerationResponse(BaseModel):
    """OpenAI compatible image generation response"""
    created: int = Field(..., description="Unix timestamp of when the image was created")
    data: List[ImageData] = Field(..., description="List of generated images")


class ErrorResponse(BaseModel):
    """Error response format"""
    error: dict = Field(..., description="Error details")


class ModelInfo(BaseModel):
    """Model information"""
    id: str
    object: str = "model"
    created: int
    owned_by: str


class ModelsResponse(BaseModel):
    """Models list response"""
    object: str = "list"
    data: List[ModelInfo]