from typing import Optional

from pydantic import BaseModel, Field


class ClassificationResult(BaseModel):
    class_name: str
    probability: float


class ClassificationResponse(BaseModel):
    data: list[ClassificationResult]


class RemoveBackgroundRequest(BaseModel):
    image_url: str = Field(..., description='URL of image to remove background from')


class RemoveBackgroundResponse(BaseModel):
    original: Optional[str] = None
    high_resolution: Optional[str] = None
    preview: Optional[str] = None
    url: Optional[str] = None


class IconRequest(BaseModel):
    prompt: str = Field(..., min_length=1, max_length=4000)
    color: Optional[str] = None
    shape: Optional[str] = None
    style: Optional[str] = None