"""Pydantic models for the OpenAI images contract."""
import base64
import binascii

from pydantic import BaseModel, field_validator

DEFAULT_SIZE = "1024x1024"


class ImageGenerationRequest(BaseModel):
    prompt: str
    model: str | None = None
    n: int | None = 1
    size: str = DEFAULT_SIZE
    seed: int | None = None
    response_format: str | None = None  # accepted; only b64_json is produced

    @field_validator("size")
    @classmethod
    def _size_shape(cls, v: str) -> str:
        w_str, sep, h_str = v.lower().partition("x")
        if not sep or not w_str.isdigit() or not h_str.isdigit() or int(w_str) <= 0 or int(h_str) <= 0:
            raise ValueError(f"invalid size '{v}', expected format 'WIDTHxHEIGHT'")
        return v


class ImageEditJsonRequest(BaseModel):
    image: str  # base64-encoded PNG/JPEG (data-URL prefix allowed)
    prompt: str
    model: str | None = None
    size: str | None = None  # accepted but ignored: edits size comes from the input image
    seed: int | None = None

    @field_validator("image")
    @classmethod
    def _decode_b64(cls, v: str) -> str:
        raw = v.split(",", 1)[1] if v.startswith("data:") else v
        try:
            base64.b64decode(raw, validate=True)
        except (binascii.Error, ValueError) as e:
            raise ValueError("image must be valid base64") from e
        return raw


class ImageItem(BaseModel):
    b64_json: str


class ImagesResponse(BaseModel):
    created: int
    data: list[ImageItem]


class ModelCard(BaseModel):
    id: str
    object: str = "model"
    owned_by: str = "comfy-adapter"


class ModelsResponse(BaseModel):
    object: str = "list"
    data: list[ModelCard]