"""图生图 / 图像编辑请求响应模型（多模型）。"""
from __future__ import annotations

from typing import Literal

from pydantic import BaseModel, Field, field_validator

# 支持的模型白名单
SUPPORTED_MODELS: frozenset[str] = frozenset(
    {
        "wanx2.1-imageedit",
        "wan2.6-image",
        "qwen-image-edit",
        "qwen-image-edit-plus",
        "qwen-image-edit-plus-2025-10-30",
    }
)

# 万相 2.1 通用图像编辑支持的 function 白名单
# description_edit             ：文本指令编辑（无掩码，即“图生图”）
# description_edit_with_mask   ：局部重绘（需 mask）
# stylization_local            ：局部风格化
# stylization_all              ：整体风格化
ImageEditFunction = Literal[
    "description_edit",
    "description_edit_with_mask",
    "stylization_local",
    "stylization_all",
]

# 单张 base64 图片（编码后）上限，超出直接拒绝，避免无谓占用内存。
_MAX_BASE64_LENGTH = 15 * 1024 * 1024


class ImageEditRequest(BaseModel):
    """图生图请求。图片以 data URL（data:image/...;base64,...）或公网 URL 传入。"""

    base_image: str = Field(..., description="待编辑原图：data URL 或公网 http(s) URL")
    prompt: str = Field(..., description="编辑指令，例如“把背景换成纯白色摄影棚”")
    model: str = Field(
        default="wanx2.1-imageedit",
        description="模型 code，见 SUPPORTED_MODELS",
    )
    mask_image: str | None = Field(
        default=None, description="掩码图（仅 function=description_edit_with_mask 时使用）"
    )
    function: ImageEditFunction = Field(
        default="description_edit", description="编辑能力类型，见 ImageEditFunction"
    )
    n: int = Field(default=1, ge=1, le=4, description="生成数量")
    size: str = Field(default="", description="输出尺寸，如 1024*1024；留空由模型决定")
    seed: int | None = Field(default=None, ge=0, description="随机种子，可选")
    style: str = Field(default="", description="输出风格，如 auto/photography/anime 等，可选")
    prompt_extend: bool = Field(default=True, description="是否自动扩写提示词（官方示例默认开启）")
    strength: float | None = Field(
        default=None,
        ge=0.0,
        le=1.0,
        description=(
            "修改幅度 0.0~1.0，官方默认 0.5；值越大改动越明显、越偏离原图。"
            "新增文字、大幅修改等场景建议 0.8 左右"
        ),
    )

    @field_validator("prompt")
    @classmethod
    def prompt_not_empty(cls, value: str) -> str:
        text = (value or "").strip()
        if not text:
            raise ValueError("prompt 不能为空")
        return text

    @field_validator("model")
    @classmethod
    def model_whitelist(cls, value: str) -> str:
        model = (value or "").strip()
        if model not in SUPPORTED_MODELS:
            raise ValueError(f"不支持的模型「{model}」，可选：{sorted(SUPPORTED_MODELS)}")
        return model

    @field_validator("base_image")
    @classmethod
    def base_image_valid(cls, value: str) -> str:
        image = (value or "").strip()
        if not image:
            raise ValueError("base_image 不能为空")
        if image.startswith("data:image/"):
            if len(image) > _MAX_BASE64_LENGTH:
                raise ValueError("图片过大，请压缩后重试")
        elif not image.startswith(("http://", "https://")):
            raise ValueError("base_image 仅支持 data URL 或公网 http(s) URL")
        return image

    @field_validator("mask_image")
    @classmethod
    def mask_image_valid(cls, value: str | None) -> str | None:
        if value is None or not value.strip():
            return None
        image = value.strip()
        if image.startswith("data:image/"):
            if len(image) > _MAX_BASE64_LENGTH:
                raise ValueError("mask_image 图片过大")
        elif not image.startswith(("http://", "https://")):
            raise ValueError("mask_image 仅支持 data URL 或公网 http(s) URL")
        return image

    @field_validator("style")
    @classmethod
    def normalize_style(cls, value: str) -> str:
        return (value or "").strip()


class ImageEditResult(BaseModel):
    url: str = Field(..., description="生成图片 URL（有效期 24 小时）")
    image_base64: str = Field(
        default="",
        description="生成图片的 data URL（服务端代理下载，便于前端直接合成/导出，规避跨域画布污染）",
    )


class ImageEditResponse(BaseModel):
    task_id: str = ""
    results: list[ImageEditResult] = Field(default_factory=list)
    image_count: int = 0
    request_id: str = ""