from pydantic import BaseModel, Field
from typing import Optional, List, Dict, Any

# AI生成相关模型
class AIGenerationRequest(BaseModel):
    prompt: str = Field(..., description="生成图片的提示词")
    negative_prompt: Optional[str] = Field(None, description="反向提示词")
    size: str = Field("1024*1024", description="图片尺寸")
    n: int = Field(1, description="生成数量", ge=1, le=4)
    model: str = Field("wan2.6-t2i", description="使用的模型")

class AITemplateGenerationRequest(BaseModel):
    template_id: str = Field(..., description="提示词模板ID")
    params: Dict[str, str] = Field(default_factory=dict, description="提示词参数")
    negative_prompt: Optional[str] = Field(None, description="反向提示词")
    size: str = Field("1024*1024", description="图片尺寸")
    n: int = Field(1, description="生成数量", ge=1, le=4)
    model: str = Field("wan2.6-t2i", description="使用的模型")

class AITaskResponse(BaseModel):
    task_id: Optional[str] = Field(None, description="任务ID (异步任务时存在)")
    request_id: Optional[str] = Field(None, description="请求ID")
    status: Optional[str] = Field(None, description="任务状态")
    results: Optional[List[Dict[str, Any]]] = Field(None, description="同步生成的直接结果")
    # results结构: [{"url": "/media/xxx.png", "origin_url": "https://..."}]

class AITaskResult(BaseModel):
    task_id: str
    status: str
    code: Optional[str] = None
    message: Optional[str] = None
    output: Optional[Dict[str, Any]] = None
    usage: Optional[Dict[str, Any]] = None