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, 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


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


class GenerationResponse(BaseModel):
    id: str
    status: GenerationStatus
    gen_type: GenType = GenType.IMAGE
    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
    # Video-specific
    kling_task_id: Optional[str] = None
    video_duration: Optional[int] = None
    video_mode: 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