# services/comfy-agent/app/models.py
from pydantic import BaseModel, Field
from typing import Any, Dict, Optional, Literal

GenType = Literal["text-to-image", "text-to-video", "image-to-video"]

class GenerateImageRequest(BaseModel):
    prompt: str = Field(min_length=1)
    negative_prompt: Optional[str] = None
    width: int = 1024
    height: int = 1024
    steps: int = 28
    seed: Optional[int] = None
    idempotency_key: Optional[str] = None
    workflow: Optional[str] = None
    workflow_params: Dict[str, Any] = Field(default_factory=dict)

class GenerateVideoRequest(BaseModel):
    prompt: str = Field(min_length=1)
    negative_prompt: Optional[str] = None
    width: int = 768
    height: int = 512
    frames: Optional[int] = None
    seconds: int = 4
    fps: int = 24
    steps: int = 30
    cfg: float = 2.5
    seed: Optional[int] = None
    format: str = "mp4"
    codec: str = "h264"
    idempotency_key: Optional[str] = None
    workflow: Optional[str] = None
    workflow_params: Dict[str, Any] = Field(default_factory=dict)

class JobStatus(BaseModel):
    job_id: str
    type: GenType
    status: Literal["queued", "running", "succeeded", "failed"]
    progress: float = 0.0
    message: Optional[str] = None
    result_url: Optional[str] = None
    error: Optional[str] = None
    comfy_prompt_id: Optional[str] = None