"""Types for Kling video and image generation."""
from __future__ import annotations

from enum import Enum
from pydantic import BaseModel, Field


class KlingVideoModel(str, Enum):
    """Available Kling video models."""
    KLING_2_6 = "kling_2.6"
    KLING_2_6_MOTION = "kling_2.6_motion"
    KLING_2_6_MOTION_PRO = "kling_2.6_motion_pro"
    KLING_3_MOTION = "kling_3_motion"
    KLING_3_MOTION_PRO = "kling_3_motion_pro"
    KLING_2_5_TURBO = "kling_2.5_turbo"
    KLING_2_5_TURBO_PRO = "kling_2.5_turbo_pro"
    KLING_2_1_MASTER = "kling_2.1_master"
    KLING_2_1 = "kling_2.1"
    KLING_2_1_PRO = "kling_2.1_pro"
    KLING_1_6 = "kling_1.6"
    KLING_1_6_PRO = "kling_1.6_pro"


class KlingO3VideoModel(str, Enum):
    """Kling O3 video models."""
    STANDARD = "kling_o3"
    PRO = "kling_o3_pro"


class KlingVideoDuration(str, Enum):
    """Video duration options."""
    FIVE = "5"
    TEN = "10"


class KlingO3VideoDuration(str, Enum):
    """O3 video duration options (supports 15s)."""
    FIVE = "5"
    TEN = "10"
    FIFTEEN = "15"


class KlingVideoAspectRatio(str, Enum):
    """Video aspect ratios."""
    LANDSCAPE = "16:9"
    PORTRAIT = "9:16"
    SQUARE = "1:1"


class KlingMotionDirection(str, Enum):
    """Motion control reference direction."""
    VIDEO = "video"
    IMAGE = "image"


class KlingPhotoResolution(str, Enum):
    """Photo resolution options."""
    RES_1K = "1k"
    RES_2K = "2k"
    RES_4K = "4k"


class KlingPhotoAspectRatio(str, Enum):
    """Photo aspect ratios."""
    AUTO = "auto"
    PORTRAIT_916 = "9:16"
    PORTRAIT_23 = "2:3"
    PORTRAIT_34 = "3:4"
    SQUARE = "1:1"
    LANDSCAPE_43 = "4:3"
    LANDSCAPE_32 = "3:2"
    LANDSCAPE_169 = "16:9"
    ULTRA_WIDE = "21:9"


class KlingStatus(str, Enum):
    """Generation status."""
    PENDING = "pending"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"


class KlingVideoGenerateParams(BaseModel):
    """Video generation request parameters."""
    prompt: str = Field(..., min_length=1, max_length=2500)
    negative_prompt: str | None = Field(default=None, max_length=2500)
    model: KlingVideoModel = Field(default=KlingVideoModel.KLING_2_5_TURBO)
    duration: KlingVideoDuration = Field(default=KlingVideoDuration.FIVE)
    aspect_ratio: KlingVideoAspectRatio = Field(default=KlingVideoAspectRatio.LANDSCAPE)
    cfg_scale: float = Field(default=0.5, ge=0.0, le=1.0)
    seed: int | None = Field(default=None, ge=0, le=2147483647)
    callback_url: str | None = None
    image_url: str | None = None
    image_base64: str | None = None
    reference_video_url: str | None = None
    reference_video_base64: str | None = None
    reference_image_url: str | None = None
    reference_image_base64: str | None = None
    motion_direction: KlingMotionDirection = Field(default=KlingMotionDirection.VIDEO)

    def to_api_body(self) -> dict:
        data: dict = {
            "prompt": self.prompt,
            "model": self.model.value,
            "duration": self.duration.value,
            "aspect_ratio": self.aspect_ratio.value,
            "cfg_scale": self.cfg_scale,
        }
        if self.negative_prompt:
            data["negative_prompt"] = self.negative_prompt
        if self.seed is not None:
            data["seed"] = self.seed
        if self.callback_url:
            data["callback_url"] = self.callback_url
        if self.image_url:
            data["image_url"] = self.image_url
        if self.image_base64:
            data["image_base64"] = self.image_base64
        if self.reference_video_url:
            data["reference_video_url"] = self.reference_video_url
        if self.reference_video_base64:
            data["reference_video_base64"] = self.reference_video_base64
        if self.reference_image_url:
            data["reference_image_url"] = self.reference_image_url
        if self.reference_image_base64:
            data["reference_image_base64"] = self.reference_image_base64
        if self.reference_video_url or self.reference_video_base64:
            data["motion_direction"] = self.motion_direction.value
        return data


class KlingO1VideoGenerateParams(BaseModel):
    """O1 video generation request parameters."""
    prompt: str = Field(..., min_length=1, max_length=2500)
    duration: KlingVideoDuration = Field(default=KlingVideoDuration.FIVE)
    aspect_ratio: KlingVideoAspectRatio | None = None
    keep_original_sound: bool = True
    callback_url: str | None = None
    images_urls: list[str] | None = Field(default=None, max_length=5)
    images_base64: list[str] | None = Field(default=None, max_length=5)
    first_frame_image_url: str | None = None
    first_frame_image_base64: str | None = None
    end_frame_image_url: str | None = None
    end_frame_image_base64: str | None = None

    def to_api_body(self) -> dict:
        data: dict = {
            "prompt": self.prompt,
            "duration": self.duration.value,
            "keep_original_sound": self.keep_original_sound,
        }
        if self.aspect_ratio:
            data["aspect_ratio"] = self.aspect_ratio.value
        if self.callback_url:
            data["callback_url"] = self.callback_url
        if self.images_urls:
            data["images_urls"] = self.images_urls
        if self.images_base64:
            data["images_base64"] = self.images_base64
        if self.first_frame_image_url:
            data["first_frame_image_url"] = self.first_frame_image_url
        if self.first_frame_image_base64:
            data["first_frame_image_base64"] = self.first_frame_image_base64
        if self.end_frame_image_url:
            data["end_frame_image_url"] = self.end_frame_image_url
        if self.end_frame_image_base64:
            data["end_frame_image_base64"] = self.end_frame_image_base64
        return data


class KlingO3VideoGenerateParams(BaseModel):
    """O3 video generation request parameters."""
    prompt: str = Field(..., min_length=1, max_length=2500)
    model: KlingO3VideoModel = Field(default=KlingO3VideoModel.STANDARD)
    duration: KlingO3VideoDuration = Field(default=KlingO3VideoDuration.FIVE)
    aspect_ratio: KlingVideoAspectRatio = Field(default=KlingVideoAspectRatio.LANDSCAPE)
    keep_original_sound: bool = True
    callback_url: str | None = None
    images_url: list[str] | None = Field(default=None, max_length=5)
    images_base64: list[str] | None = Field(default=None, max_length=5)
    first_frame_url: str | None = None
    first_frame_base64: str | None = None
    end_frame_url: str | None = None
    end_frame_base64: str | None = None
    customize_multi_shots: bool = False
    prefer_multi_shots: bool = False

    def to_api_body(self) -> dict:
        data: dict = {
            "prompt": self.prompt,
            "model": self.model.value,
            "duration": self.duration.value,
            "aspect_ratio": self.aspect_ratio.value,
            "keep_original_sound": self.keep_original_sound,
        }
        if self.callback_url:
            data["callback_url"] = self.callback_url
        if self.images_url:
            data["images_url"] = self.images_url
        if self.images_base64:
            data["images_base64"] = self.images_base64
        if self.first_frame_url:
            data["first_frame_url"] = self.first_frame_url
        if self.first_frame_base64:
            data["first_frame_base64"] = self.first_frame_base64
        if self.end_frame_url:
            data["end_frame_url"] = self.end_frame_url
        if self.end_frame_base64:
            data["end_frame_base64"] = self.end_frame_base64
        if self.customize_multi_shots:
            data["customize_multi_shots"] = True
        if self.prefer_multi_shots:
            data["prefer_multi_shots"] = True
        return data


class KlingO1PhotoGenerateParams(BaseModel):
    """O1 photo generation request parameters."""
    prompt: str = Field(..., min_length=1, max_length=2500)
    img_resolution: KlingPhotoResolution = Field(default=KlingPhotoResolution.RES_2K)
    aspect_ratio: KlingPhotoAspectRatio = Field(default=KlingPhotoAspectRatio.SQUARE)
    callback_url: str | None = None
    images_url: list[str] | None = Field(default=None, max_length=10)
    images_base64: list[str] | None = Field(default=None, max_length=10)

    def to_api_body(self) -> dict:
        data: dict = {
            "prompt": self.prompt,
            "img_resolution": self.img_resolution.value,
            "aspect_ratio": self.aspect_ratio.value,
        }
        if self.callback_url:
            data["callback_url"] = self.callback_url
        if self.images_url:
            data["images_url"] = self.images_url
        if self.images_base64:
            data["images_base64"] = self.images_base64
        return data


class KlingO3PhotoGenerateParams(BaseModel):
    """O3 photo generation request parameters."""
    prompt: str = Field(..., min_length=1, max_length=2500)
    img_resolution: KlingPhotoResolution = Field(default=KlingPhotoResolution.RES_4K)
    aspect_ratio: KlingPhotoAspectRatio = Field(default=KlingPhotoAspectRatio.SQUARE)
    prefer_multi_shots: bool = False
    callback_url: str | None = None
    images_url: list[str] | None = Field(default=None, max_length=10)
    images_base64: list[str] | None = Field(default=None, max_length=10)

    def to_api_body(self) -> dict:
        data: dict = {
            "prompt": self.prompt,
            "img_resolution": self.img_resolution.value,
            "aspect_ratio": self.aspect_ratio.value,
        }
        if self.prefer_multi_shots:
            data["prefer_multi_shots"] = True
        if self.callback_url:
            data["callback_url"] = self.callback_url
        if self.images_url:
            data["images_url"] = self.images_url
        if self.images_base64:
            data["images_base64"] = self.images_base64
        return data


class KlingGenerateResponse(BaseModel):
    """Response from video generation endpoints."""
    generation_id: str
    status: KlingStatus
    price: str

class KlingVideoStatusResponse(BaseModel):
    """Video generation status response."""
    generation_id: str
    status: KlingStatus
    prompt: str
    video_url: str | None = None
    model: str | None = None
    description: str | None = None

    @property
    def is_terminal(self) -> bool:
        return self.status in (KlingStatus.COMPLETED, KlingStatus.FAILED)

    @property
    def is_completed(self) -> bool:
        return self.status == KlingStatus.COMPLETED

    @property
    def is_failed(self) -> bool:
        return self.status == KlingStatus.FAILED


class KlingPhotoGenerateResponse(BaseModel):
    """Response from photo generation endpoints."""
    generation_id: str
    status: KlingStatus
    price: str


class KlingPhotoStatusResponse(BaseModel):
    """Photo generation status response."""
    generation_id: str
    status: KlingStatus
    prompt: str
    image_url: str | None = None
    model: str | None = None
    description: str | None = None

    @property
    def is_terminal(self) -> bool:
        return self.status in (KlingStatus.COMPLETED, KlingStatus.FAILED)

    @property
    def is_completed(self) -> bool:
        return self.status == KlingStatus.COMPLETED

    @property
    def is_failed(self) -> bool:
        return self.status == KlingStatus.FAILED
