from datetime import datetime, UTC
from typing import List, Optional

from pydantic import BaseModel

from models.Asset import Asset
from models.Generation import GenerationStatus
from models.enums import AspectRatios, Quality


class GenerationRequest(BaseModel):
    linked_character_id: Optional[str] = None
    aspect_ratio: AspectRatios = AspectRatios.NINESIXTEEN
    quality: Quality = Quality.ONEK
    prompt: str
    assets_list: List[str]


class GenerationResponse(BaseModel):
    id: str
    status: GenerationStatus
    linked_character_id: Optional[str] = None
    aspect_ratio: AspectRatios
    quality: Quality
    prompt: str
    assets_list: List[str]
    result: Optional[str] = None
    created_at: datetime = datetime.now(UTC)
    updated_at: datetime = datetime.now(UTC)



class PromptRequest(BaseModel):
    prompt: str
    linked_assets: List[str] = []


class PromptResponse(BaseModel):
    prompt: str