from typing import Optional
from pydantic import BaseModel, Field
from models.enums import AspectRatios, Quality


class ExternalGenerationRequest(BaseModel):
    """Request model for importing external generations."""
    
    prompt: str
    tech_prompt: Optional[str] = None
    
    # Image can be provided as base64 string OR URL (one must be provided)
    image_data: Optional[str] = Field(None, description="Base64-encoded image data")
    image_url: Optional[str] = Field(None, description="URL to download image from")

    nsfw: bool = False

    # Generation metadata
    aspect_ratio: AspectRatios = AspectRatios.NINESIXTEEN  # "1:1","2:3","3:2","3:4","4:3","4:5","5:4","9:16","16:9","21:9"
    quality: Quality = Quality.ONEK
    
    # Optional linking
    linked_character_id: Optional[str] = None
    created_by: str = Field(..., description="User ID from external system")
    project_id: Optional[str] = None
    
    # Performance metrics
    execution_time_seconds: Optional[float] = None
    api_execution_time_seconds: Optional[float] = None
    token_usage: Optional[int] = None
    input_token_usage: Optional[int] = None
    output_token_usage: Optional[int] = None
    
    def validate_image_source(self):
        """Ensure at least one image source is provided."""
        if not self.image_data and not self.image_url:
            raise ValueError("Either image_data or image_url must be provided")
        if self.image_data and self.image_url:
            raise ValueError("Only one of image_data or image_url should be provided")