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

from pydantic import BaseModel, Field

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


class GenerationRequest(BaseModel):
    linked_character_id: Optional[str] = None
    aspect_ratio: AspectRatios = AspectRatios.NINESIXTEEN
    quality: Quality = Quality.ONEK
    prompt: str
    telegram_id: Optional[int] = None
    use_profile_image: bool = True
    assets_list: List[str]
    project_id: Optional[str] = None
    idea_id: Optional[str] = None
    count: int = Field(default=1, ge=1, le=10)


class GenerationsResponse(BaseModel):
    generations: List["GenerationResponse"]
    total_count: int


class GenerationResponse(BaseModel):
    id: str
    status: GenerationStatus
    failed_reason: Optional[str] = None
    
    linked_character_id: Optional[str] = None
    aspect_ratio: AspectRatios
    quality: Quality
    prompt: str
    tech_prompt: Optional[str] = None
    assets_list: List[str]
    result_list: List[str] = []
    result: Optional[str] = None
    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
    progress: int = 0
    cost: Optional[float] = None
    created_by: Optional[str] = None
    generation_group_id: Optional[str] = None
    idea_id: Optional[str] = None
    created_at: datetime = datetime.now(UTC)
    updated_at: datetime = datetime.now(UTC)


class GenerationGroupResponse(BaseModel):
    generation_group_id: str
    generations: List[GenerationResponse]


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


class PromptResponse(BaseModel):
    prompt: str