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 GenerateProductImageModel(BaseModel):
    tasks_id: str
    prompt: str
    image_url: str
    image_strength: float
    product_type: str


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]