freepik/freepik_cli/api/models.py

"""Enums, endpoint maps, and response normalization utilities."""

from __future__ import annotations

from enum import Enum
from typing import Any


# ---------------------------------------------------------------------------
# Model enums
# ---------------------------------------------------------------------------

class ImageModel(str, Enum):
    MYSTIC = "mystic"
    FLUX_KONTEXT_PRO = "flux-kontext-pro"
    FLUX_2_PRO = "flux-2-pro"
    FLUX_2_TURBO = "flux-2-turbo"
    FLUX_PRO_1_1 = "flux-pro-1.1"
    SEEDREAM_V4 = "seedream-v4"
    SEEDREAM_V4_5 = "seedream-v4-5"
    IDEOGRAM_V2 = "ideogram-v2"


class VideoModel(str, Enum):
    KLING_O1_PRO = "kling-o1-pro"
    KLING_O1_STD = "kling-o1-std"
    KLING_ELEMENTS_PRO = "kling-elements-pro"
    KLING_ELEMENTS_STD = "kling-elements-std"
    MINIMAX_HAILUO = "minimax-hailuo"


class UpscaleMode(str, Enum):
    CREATIVE = "creative"
    PRECISION = "precision"
    PRECISION_V2 = "precision-v2"


class VideoUpscaleMode(str, Enum):
    STANDARD = "standard"
    TURBO = "turbo"


class AspectRatio(str, Enum):
    LANDSCAPE = "16:9"
    PORTRAIT = "9:16"
    SQUARE = "1:1"
    CLASSIC = "4:3"
    WIDE = "21:9"


class IconStyle(str, Enum):
    SOLID = "solid"
    OUTLINE = "outline"
    COLOR = "color"
    FLAT = "flat"
    STICKER = "sticker"


# ---------------------------------------------------------------------------
# Endpoint maps
# ---------------------------------------------------------------------------

IMAGE_POST_ENDPOINTS: dict[ImageModel, str] = {
    ImageModel.MYSTIC: "/v1/ai/mystic",
    ImageModel.FLUX_KONTEXT_PRO: "/v1/ai/text-to-image/flux-kontext-pro",
    ImageModel.FLUX_2_PRO: "/v1/ai/text-to-image/flux-2-pro",
    ImageModel.FLUX_2_TURBO: "/v1/ai/text-to-image/flux-2-turbo",
    ImageModel.FLUX_PRO_1_1: "/v1/ai/text-to-image/flux-pro-v1-1",
    ImageModel.SEEDREAM_V4: "/v1/ai/text-to-image/seedream-v4",
    ImageModel.SEEDREAM_V4_5: "/v1/ai/text-to-image/seedream-v4-5",
    ImageModel.IDEOGRAM_V2: "/v1/ai/text-to-image/ideogram-v2",
}

IMAGE_STATUS_ENDPOINTS: dict[ImageModel, str] = {
    ImageModel.MYSTIC: "/v1/ai/mystic/{task_id}",
    ImageModel.FLUX_KONTEXT_PRO: "/v1/ai/text-to-image/flux-kontext-pro/{task_id}",
    ImageModel.FLUX_2_PRO: "/v1/ai/text-to-image/flux-2-pro/{task_id}",
    ImageModel.FLUX_2_TURBO: "/v1/ai/text-to-image/flux-2-turbo/{task_id}",
    ImageModel.FLUX_PRO_1_1: "/v1/ai/text-to-image/flux-pro-v1-1/{task_id}",
    ImageModel.SEEDREAM_V4: "/v1/ai/text-to-image/seedream-v4/{task_id}",
    ImageModel.SEEDREAM_V4_5: "/v1/ai/text-to-image/seedream-v4-5/{task_id}",
    ImageModel.IDEOGRAM_V2: "/v1/ai/text-to-image/ideogram-v2/{task_id}",
}

VIDEO_POST_ENDPOINTS: dict[VideoModel, str] = {
    VideoModel.KLING_O1_PRO: "/v1/ai/image-to-video/kling-o1-pro",
    VideoModel.KLING_O1_STD: "/v1/ai/image-to-video/kling-o1-std",
    VideoModel.KLING_ELEMENTS_PRO: "/v1/ai/image-to-video/kling-elements-pro",
    VideoModel.KLING_ELEMENTS_STD: "/v1/ai/image-to-video/kling-elements-std",
    VideoModel.MINIMAX_HAILUO: "/v1/ai/image-to-video/minimax-hailuo-02-1080p",
}

VIDEO_STATUS_ENDPOINTS: dict[VideoModel, str] = {
    VideoModel.KLING_O1_PRO: "/v1/ai/image-to-video/kling-o1/{task_id}",
    VideoModel.KLING_O1_STD: "/v1/ai/image-to-video/kling-o1/{task_id}",
    VideoModel.KLING_ELEMENTS_PRO: "/v1/ai/image-to-video/kling-elements-pro/{task_id}",
    VideoModel.KLING_ELEMENTS_STD: "/v1/ai/image-to-video/kling-elements-std/{task_id}",
    VideoModel.MINIMAX_HAILUO: "/v1/ai/image-to-video/minimax-hailuo-02-1080p/{task_id}",
}

# Per-model image input field name
VIDEO_IMAGE_FIELDS: dict[VideoModel, str] = {
    VideoModel.KLING_O1_PRO: "first_frame",
    VideoModel.KLING_O1_STD: "first_frame",
    VideoModel.KLING_ELEMENTS_PRO: "images",    # expects an array
    VideoModel.KLING_ELEMENTS_STD: "images",    # expects an array
    VideoModel.MINIMAX_HAILUO: "image",
}

# Kling Elements requires named aspect ratio slugs
VIDEO_SLUG_ASPECT_RATIO_MODELS: set[VideoModel] = {
    VideoModel.KLING_ELEMENTS_PRO,
    VideoModel.KLING_ELEMENTS_STD,
}

_VIDEO_RATIO_TO_SLUG: dict[str, str] = {
    "1:1":  "square_1_1",
    "16:9": "widescreen_16_9",
    "9:16": "social_story_9_16",
}


def normalize_aspect_ratio_video(ratio: str, model: VideoModel) -> str:
    """Convert a user-facing aspect ratio to the format required by the video model."""
    if model not in VIDEO_SLUG_ASPECT_RATIO_MODELS:
        return ratio
    return _VIDEO_RATIO_TO_SLUG.get(ratio, ratio)

UPSCALE_POST_ENDPOINTS: dict[UpscaleMode, str] = {
    UpscaleMode.CREATIVE: "/v1/ai/image-upscaler",
    UpscaleMode.PRECISION: "/v1/ai/image-upscaler-precision",
    UpscaleMode.PRECISION_V2: "/v1/ai/image-upscaler-precision-v2",
}

UPSCALE_STATUS_ENDPOINTS: dict[UpscaleMode, str] = {
    UpscaleMode.CREATIVE: "/v1/ai/image-upscaler/{task_id}",
    UpscaleMode.PRECISION: "/v1/ai/image-upscaler-precision/{task_id}",
    UpscaleMode.PRECISION_V2: "/v1/ai/image-upscaler-precision-v2/{task_id}",
}

VIDEO_UPSCALE_POST_ENDPOINTS: dict[VideoUpscaleMode, str] = {
    VideoUpscaleMode.STANDARD: "/v1/ai/video-upscaler",
    VideoUpscaleMode.TURBO: "/v1/ai/video-upscaler/turbo",
}

VIDEO_UPSCALE_STATUS_ENDPOINT = "/v1/ai/video-upscaler/{task_id}"


# ---------------------------------------------------------------------------
# Aspect ratio normalization
# ---------------------------------------------------------------------------

# Models that require named slug aspect ratios instead of "W:H" strings
SLUG_ASPECT_RATIO_MODELS: set[ImageModel] = {
    ImageModel.MYSTIC,
    ImageModel.FLUX_PRO_1_1,
    ImageModel.SEEDREAM_V4,
    ImageModel.SEEDREAM_V4_5,
}

# User-friendly "W:H" → API slug mapping
_RATIO_TO_SLUG: dict[str, str] = {
    "1:1":   "square_1_1",
    "16:9":  "widescreen_16_9",
    "9:16":  "social_story_9_16",
    "4:3":   "classic_4_3",
    "3:4":   "traditional_3_4",
    "3:2":   "standard_3_2",
    "2:3":   "portrait_2_3",
    "2:1":   "horizontal_2_1",
    "1:2":   "vertical_1_2",
    "4:5":   "social_post_4_5",
    "21:9":  "widescreen_16_9",   # closest match
}


def normalize_aspect_ratio(ratio: str, model: ImageModel) -> str:
    """Convert a user-facing aspect ratio to the format required by the model."""
    if model not in SLUG_ASPECT_RATIO_MODELS:
        return ratio  # free-form models accept "1:1" directly
    slug = _RATIO_TO_SLUG.get(ratio)
    if slug:
        return slug
    # Already a slug (user passed "square_1_1" directly) — pass through
    return ratio


# ---------------------------------------------------------------------------
# Response normalization helpers
# ---------------------------------------------------------------------------

def get_task_id(raw: dict[str, Any]) -> str:
    """Extract task_id from any response shape."""
    data = raw.get("data", raw)
    task_id = data.get("task_id") or data.get("id") or raw.get("task_id") or raw.get("id")
    if not task_id:
        raise ValueError(f"No task_id found in response: {raw}")
    return str(task_id)


def get_status(raw: dict[str, Any]) -> str:
    """Extract normalized status string from any response shape."""
    data = raw.get("data", raw)
    status = data.get("status") or raw.get("status") or "PENDING"
    return status.upper().replace(" ", "_")


def get_output_urls(raw: dict[str, Any]) -> list[str]:
    """Extract all output file URLs from a completed task response."""
    data = raw.get("data", raw)

    # Try common key names in order of likelihood
    for key in ("generated", "output", "outputs", "result", "results", "images", "videos"):
        items = data.get(key)
        if items is None:
            continue
        if isinstance(items, list) and items:
            urls = []
            for item in items:
                if isinstance(item, dict):
                    url = item.get("url") or item.get("download_url") or item.get("src")
                    if url:
                        urls.append(url)
                elif isinstance(item, str):
                    urls.append(item)
            if urls:
                return urls
        elif isinstance(items, dict):
            url = items.get("url") or items.get("download_url") or items.get("src")
            if url:
                return [url]
        elif isinstance(items, str):
            return [items]

    # Fallback: top-level url field
    if "url" in data:
        return [data["url"]]

    return []