from datetime import datetime, UTC

from pydantic import BaseModel, Field

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


class GenerationRequest(BaseModel):
    linked_character_id: str | None = None
    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
    prompt: str
    model: ImageModel = Field(default=ImageModel.GEMINI_3_PRO_IMAGE_PREVIEW)
    telegram_id: int | None = None
    use_profile_image: bool = True
    assets_list: list[str]
    environment_id: str | None = None
    project_id: str | None = None
    idea_id: str | None = None
    nsfw: bool = False
    count: int = Field(default=1, ge=1, le=10)


class NsfwRequest(BaseModel):
    is_nsfw: bool


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


class GenerationResponse(BaseModel):
    id: str
    status: GenerationStatus
    failed_reason: str | None = None
    project_id: str | None = None
    linked_character_id: str | None = None
    aspect_ratio: AspectRatios
    quality: Quality
    prompt: str
    model: ImageModel | None = None
    seed: int | None = None
    tech_prompt: str | None = None
    assets_list: list[str]
    result_list: list[str] = []
    result: str | None = None
    execution_time_seconds: float | None = None
    api_execution_time_seconds: float | None = None
    token_usage: int | None = None
    input_token_usage: int | None = None
    output_token_usage: int | None = None
    progress: int = 0
    cost: float | None = None
    created_by: str | None = None
    generation_group_id: str | None = None
    idea_id: str | None = None
    likes_count: int = 0
    is_liked: bool = False
    nsfw: bool = False
    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
    model: TextModel = Field(default=TextModel.GEMINI_3_1_PRO_PREVIEW)
    linked_assets: list[str] = []


class PromptResponse(BaseModel):
    prompt: str