from __future__ import annotations

from enum import Enum
from typing import Optional, Union

from pydantic import BaseModel, Field, confloat


class LumaImageModel(str, Enum):
    photon_1 = "photon-1"
    photon_flash_1 = "photon-flash-1"


class LumaVideoModel(str, Enum):
    ray_2 = "ray-2"
    ray_flash_2 = "ray-flash-2"
    ray_1_6 = "ray-1-6"


class LumaAspectRatio(str, Enum):
    ratio_1_1 = "1:1"
    ratio_16_9 = "16:9"
    ratio_9_16 = "9:16"
    ratio_4_3 = "4:3"
    ratio_3_4 = "3:4"
    ratio_21_9 = "21:9"
    ratio_9_21 = "9:21"


class LumaVideoOutputResolution(str, Enum):
    res_540p = "540p"
    res_720p = "720p"
    res_1080p = "1080p"
    res_4k = "4k"


class LumaVideoModelOutputDuration(str, Enum):
    dur_5s = "5s"
    dur_9s = "9s"


class LumaGenerationType(str, Enum):
    video = 'video'
    image = 'image'


class LumaState(str, Enum):
    queued = "queued"
    dreaming = "dreaming"
    completed = "completed"
    failed = "failed"


class LumaAssets(BaseModel):
    video: Optional[str] = Field(None, description='The URL of the video')
    image: Optional[str] = Field(None, description='The URL of the image')
    progress_video: Optional[str] = Field(None, description='The URL of the progress video')


class LumaImageRef(BaseModel):
    '''Used for image gen'''
    url: str = Field(..., description='The URL of the image reference')
    weight: confloat(ge=0.0, le=1.0) = Field(..., description='The weight of the image reference')


class LumaImageReference(BaseModel):
    '''Used for video gen'''
    type: str = Field('image', description='Input type, defaults to image')
    url: str = Field(..., description='The URL of the image')


class LumaModifyImageRef(BaseModel):
    url: str = Field(..., description='The URL of the image reference')
    weight: confloat(ge=0.0, le=1.0) = Field(..., description='The weight of the image reference')


class LumaCharacterRef(BaseModel):
    identity0: LumaImageIdentity = Field(..., description='The image identity object')


class LumaImageIdentity(BaseModel):
    images: list[str] = Field(..., description='The URLs of the image identity')


class LumaGenerationReference(BaseModel):
    type: str = Field('generation', description='Input type, defaults to generation')
    id: str = Field(..., description='The ID of the generation')


class LumaKeyframe(BaseModel):
    generation: Optional[LumaGenerationReference] = Field(None, description='Reference to generation')
    image: Optional[LumaImageReference] = Field(None, description='Reference to image')


class LumaKeyframes(BaseModel):
    frame0: Optional[LumaKeyframe] = Field(None, description='')
    frame1: Optional[LumaKeyframe] = Field(None, description='')


class LumaImageGenerationRequest(BaseModel):
    prompt: str = Field(..., description='The prompt of the generation')
    model: LumaImageModel = Field(LumaImageModel.photon_1, description='The image model used for the generation')
    aspect_ratio: Optional[LumaAspectRatio] = Field(LumaAspectRatio.ratio_16_9, description='The aspect ratio of the generation')
    image_ref: Optional[list[LumaImageRef]] = Field(None, description='List of image reference objects')
    style_ref: Optional[list[LumaImageRef]] = Field(None, description='List of style reference objects')
    character_ref: Optional[LumaCharacterRef] = Field(None, description='The image identity object')
    modify_image_ref: Optional[LumaModifyImageRef] = Field(None, description='The modify image reference object')


class LumaGenerationRequest(BaseModel):
    prompt: str = Field(..., description='The prompt of the generation')
    model: LumaVideoModel = Field(LumaVideoModel.ray_2, description='The video model used for the generation')
    duration: LumaVideoModelOutputDuration = Field(LumaVideoModelOutputDuration.dur_5s, description='The duration of the generation')
    aspect_ratio: Optional[LumaAspectRatio] = Field(None, description='The aspect ratio of the generation')
    resolution: Optional[LumaVideoOutputResolution] = Field(None, description='The resolution of the generation')
    loop: Optional[bool] = Field(None, description='Whether to loop the video')
    keyframes: Optional[LumaKeyframes] = Field(None, description='The keyframes of the generation')


class LumaGeneration(BaseModel):
    id: str = Field(..., description='The ID of the generation')
    generation_type: LumaGenerationType = Field(..., description='Generation type, image or video')
    state: LumaState = Field(..., description='The state of the generation')
    failure_reason: Optional[str] = Field(None, description='The reason for the state of the generation')
    created_at: str = Field(..., description='The date and time when the generation was created')
    assets: Optional[LumaAssets] = Field(None, description='The assets of the generation')
    model: str = Field(..., description='The model used for the generation')
    request: Union[LumaGenerationRequest, LumaImageGenerationRequest] = Field(..., description="The request used for the generation")