from typing import List

from pydantic import BaseModel


class GenerateMultiViewModel(BaseModel):
    tasks_id: str
    image_url: str


class GenerateImageModel(BaseModel):
    tasks_id: str
    prompt: str
    image_url: str
    mode: str
    category: str
    gender: str
    version: str


class GenerateSingleLogoImageModel(BaseModel):
    tasks_id: str
    prompt: str
    seed: str


class GenerateSloganImageModel(BaseModel):
    num_point: int
    tasks_id: str
    prompt: str
    image_url: str


class GenerateProductImageModel(BaseModel):
    tasks_id: str
    prompt: str
    image_url: str
    image_strength: float
    product_type: str


class Flux2ToProductImgModel(BaseModel):
    tasks_id: str
    prompt: str
    image_path: str
    infer_step: int | None = None


class GenerateRelightImageModel(BaseModel):
    tasks_id: str
    prompt: str
    image_url: str
    direction: str
    product_type: str


"""
    batch generate image
"""


# product任务子项
class ProductItemModel(BaseModel):
    tasks_id: str
    image_strength: float
    prompt: str
    image_url: str
    product_type: str


# product批处理 集合
class BatchGenerateProductImageModel(BaseModel):
    batch_tasks_id: str
    user_id: str
    batch_data_list: List[ProductItemModel]


# relight任务子项
class RelightItemModel(BaseModel):
    tasks_id: str
    prompt: str
    image_url: str
    direction: str
    product_type: str


# relight批处理集合
class BatchGenerateRelightImageModel(BaseModel):
    batch_tasks_id: str
    user_id: str
    batch_data_list: List[RelightItemModel]


"""
    agent tool generate image
"""


class AgentTollGenerateImageModel(BaseModel):
    prompt: str
    category: str
    gender: str
    version: str
    size: int