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

from pydantic import BaseModel, Field, computed_field

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


class GenerationStatus(str, Enum):
    RUNNING = "running"
    DONE = "done"
    FAILED = "failed"

class Generation(BaseModel):
    id: Optional[str] = None
    status: GenerationStatus = GenerationStatus.RUNNING
    failed_reason: Optional[str] = None
    linked_character_id: Optional[str] = None
    telegram_id: Optional[int] = None
    use_profile_image: bool = True
    aspect_ratio: AspectRatios
    quality: Quality
    prompt: str
    tech_prompt: Optional[str] = None
    assets_list: List[str] = Field(default_factory=list)
    result_list: List[str] = Field(default_factory=list)
    result: Optional[str] = None
    progress: int = 0
    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
    is_deleted: bool = False
    album_id: Optional[str] = None
    created_by: Optional[str] = None
    created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
    updated_at: datetime = Field(default_factory=lambda: datetime.now(UTC))

    @computed_field
    def cost(self) -> float:
        if self.status == GenerationStatus.DONE:
            cost_input = self.input_token_usage * 0.000002
            cost_output = self.output_token_usage * 0.00012
            return round(cost_input + cost_output, 3)
        return 0.0