from pydantic import BaseModel, Field, model_validator

from app.enums.generation_provider import IMAGE_MULTI_OUTPUT_MAX, IMAGE_MULTI_REFERENCE_MAX
from app.schemas.common import NaiveDatetime, NaiveDatetimeOptional


class ImageEngineCreate(BaseModel):
    name: str = Field(..., max_length=64)
    provider: str = Field(..., max_length=32)
    api_base: str = Field(..., max_length=512)
    api_key: str = Field(default="", max_length=256)
    model_name: str = Field(default="", max_length=128)
    supported_models: str = Field(default='["doubao-seedream-5-0-260128"]')
    supported_sizes: str = Field(default='{}')
    default_size: str = Field(default="2K", max_length=32)
    max_image_count: int = Field(default=0)

    multi_generation_enabled: bool = Field(
        default=False,
        description="是否允许客户端选择生成多份图片；关闭时客户端只能选择 1 份",
    )
    max_generation_count: int = Field(
        default=1,
        ge=1,
        le=5,
        description="客户端单次最多可选择的生成数量，范围 1-5",
    )
    multi_image_max_images: int = Field(
        default=IMAGE_MULTI_OUTPUT_MAX,
        ge=1,
        le=IMAGE_MULTI_OUTPUT_MAX,
        description="火山组图接口输入参考图与输出图片总上限",
    )
    max_reference_image_count: int = Field(
        default=IMAGE_MULTI_REFERENCE_MAX,
        ge=0,
        le=IMAGE_MULTI_REFERENCE_MAX,
        description="图片引擎允许的最大参考图片数量",
    )
    output_format: str = Field(
        default="",
        max_length=16,
        description="供应商输出格式；留空表示不传该参数，用于兼容不支持 output_format 的模型",
    )
    generate_url: str = Field(default="", max_length=512)
    is_active: bool = True
    priority: int = 0

    @model_validator(mode="after")
    def validate_multi_generation_capability(self):
        if self.max_generation_count > self.multi_image_max_images:
            raise ValueError("max_generation_count 不能大于 multi_image_max_images")
        normalized_output_format = (self.output_format or "").lower().strip()
        if normalized_output_format not in {"", "png", "jpeg"}:
            raise ValueError("output_format 仅支持留空、png 或 jpeg")
        self.output_format = normalized_output_format
        return self


class ImageEngineOut(ImageEngineCreate):
    id: str
    created_at: NaiveDatetime
    deleted_at: NaiveDatetimeOptional = None

    model_config = {"from_attributes": True}


class ImageEnginePublic(BaseModel):
    id: str
    name: str
    provider: str
    supported_models: list[str] = []
    supported_sizes: dict[str, dict[str, str]] = {}
    default_size: str = "2K"
    max_image_count: int = 0
    multi_generation_enabled: bool = False
    max_generation_count: int = 1
    multi_image_max_images: int = IMAGE_MULTI_OUTPUT_MAX
    max_reference_image_count: int = IMAGE_MULTI_REFERENCE_MAX


class ImageEngineListResponse(BaseModel):
    items: list[ImageEnginePublic]