models.py

from enum import Enum
from typing import Optional, Any
from datetime import datetime
from pydantic import BaseModel, Field, validator
import uuid

class TaskStatus(str, Enum):

    QUEUED = "queued"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"

class UIState(str, Enum):

    IDLE = "idle"
    SUBMITTING = "submitting"
    QUEUED = "queued"
    PROCESSING = "processing"
    COMPLETED = "completed"
    ERROR = "error"

class TaskInfo(BaseModel):

    task_id: str = Field(..., description="Unique task identifier")
    status: TaskStatus = Field(..., description="Current task status")
    position: int = Field(default=0, description="Queue position")
    message: str = Field(default="", description="Status message")
    result_url: Optional[str] = Field(default=None, description="Result image URL")
    error_message: Optional[str] = Field(default=None, description="Error message if failed")
    
    @validator('task_id')
    @classmethod
    def validate_task_id(cls, v):
    
        try:
            uuid.UUID(v)
            return v
        except ValueError:
            raise ValueError(f"Invalid task ID format: {v}")
    
    class Config:
        frozen = True  # Immutable

class UIStateInfo(BaseModel):

    state: UIState = Field(..., description="Current UI state")
    task_id: Optional[str] = Field(default=None, description="Associated task ID")
    status_text: str = Field(default="", description="Display status text")
    progress_text: str = Field(default="", description="Progress description")
    show_image: bool = Field(default=False, description="Whether to show result image")
    image_url: Optional[str] = Field(default=None, description="Result image URL")
    submit_enabled: bool = Field(default=True, description="Submit button enabled")
    cancel_enabled: bool = Field(default=False, description="Cancel button enabled")
    timestamp: datetime = Field(default_factory=datetime.now, description="State timestamp")
    
    class Config:
        frozen = True  # Immutable

class TaskSubmission(BaseModel):

    prompt: str = Field(..., min_length=5, max_length=4096, description="Image generation prompt")
    width: int = Field(default=1024, description="Image width")
    height: int = Field(default=1024, description="Image height")
    task_type: str = Field(default="text_to_image", description="Task type")

    def to_api_payload(self) -> dict[str, Any]:
    
        return {
            "task_type": self.task_type,
            "input_params": {
                "prompt": self.prompt,
                "width": self.width,
                "height": self.height
            }
        }

class ImageToImageSubmission(BaseModel):

    image_data: str = Field(..., description="Base64 encoded image data")
    task_type: str = Field(default="image_to_image", description="Task type")

    def to_api_payload(self) -> dict[str, Any]:
    
        return {
            "image_data": self.image_data
        }

class PhotoStyleSubmission(BaseModel):

    image_data: str = Field(..., description="Base64 encoded image data")
    style_preset: str = Field(..., description="Photography style preset")
    task_type: str = Field(default="photo_style_transfer", description="Task type")

    def to_api_payload(self) -> dict[str, Any]:
    
        return {
            "image_data": self.image_data,
            "style_preset": self.style_preset
        }

class InteriorDesignRenderingSubmission(BaseModel):

    image_data: str = Field(..., description="Base64 encoded white model image data")
    design_style: str = Field(..., description="Interior design style")
    task_type: str = Field(default="interior_design_rendering", description="Task type")

    def to_api_payload(self) -> dict[str, Any]:
    
        return {
            "image_data": self.image_data,
            "design_style": self.design_style
        }

class WatermarkRemovalSubmission(BaseModel):

    image_data: str = Field(..., description="Base64 encoded image data")
    task_type: str = Field(default="watermark_removal", description="Task type")

    def to_api_payload(self) -> dict[str, Any]:
    
        return {
            "image_data": self.image_data
        }

class LineArtConversionSubmission(BaseModel):

    image_data: str = Field(..., description="Base64 encoded image data")
    task_type: str = Field(default="line_art_conversion", description="Task type")

    def to_api_payload(self) -> dict[str, Any]:
    
        return {
            "image_data": self.image_data
        }

class AnimeToRealSubmission(BaseModel):

    image_data: str = Field(..., description="Base64 encoded anime image data")
    task_type: str = Field(default="anime_to_real", description="Task type")

    def to_api_payload(self) -> dict[str, Any]:
    
        return {
            "image_data": self.image_data
        }

class RealToAnimeSubmission(BaseModel):

    image_data: str = Field(..., description="Base64 encoded real image data")
    task_type: str = Field(default="real_to_anime", description="Task type")

    def to_api_payload(self) -> dict[str, Any]:
    
        return {
            "image_data": self.image_data
        }

class AppSession(BaseModel):

    session_id: str = Field(default_factory=lambda: str(uuid.uuid4()), description="Session ID")
    current_ui_state: UIStateInfo = Field(default_factory=lambda: UIStateInfo(state=UIState.IDLE), description="Current UI state")
    active_task_id: Optional[str] = Field(default=None, description="Currently active task ID")
    last_submission: Optional[TaskSubmission] = Field(default=None, description="Last task submission")
    created_at: datetime = Field(default_factory=datetime.now, description="Session creation time")
    
    def is_active_task(self, task_id: str) -> bool:
    
        return self.active_task_id == task_id
    
    def can_submit_new_task(self) -> bool:
        """Check if a new task can be submitted based on current UI state"""
        current_state = self.current_ui_state
        if hasattr(current_state, 'state'):
            return current_state.state in {UIState.IDLE, UIState.COMPLETED, UIState.ERROR}
        return True  # Default to allowing submission if state is unclear
    
    class Config:
        frozen = False  # Mutable for state updates

class ImageOutpaintingSubmission(BaseModel):

    image_data: str = Field(..., description="Base64编码的图片数据")
    expand_height: float = Field(default=0.2, description="扩展高度百分比")
    expand_width: float = Field(default=0.2, description="扩展宽度百分比")

    def to_api_payload(self) -> dict:
    
        return {
            "image_data": self.image_data,
            "expand_height": self.expand_height,
            "expand_width": self.expand_width
        }

class FiveViewGenerationSubmission(BaseModel):

    image_data: str = Field(..., description="Base64编码的图片数据")

    def to_api_payload(self) -> dict:

        return {
            "image_data": self.image_data
        }


class Figure3DSubmission(BaseModel):

    image_data: str = Field(..., description="Base64编码的2D角色图片数据")
    figure_style: str = Field(..., description="3D手办风格")
    resolution: str = Field(
        default="square - 1024x1024 (1:1)",
        description="图像分辨率"
    )

    def to_api_payload(self) -> dict:

        return {
            "image_data": self.image_data,
            "figure_style": self.figure_style,
            "resolution": self.resolution
        }


class CharacterFigureCollaborationSubmission(BaseModel):

    image_data: str = Field(..., description="Base64编码的人物全身照片数据")

    def to_api_payload(self) -> dict[str, Any]:

        return {
            "image_data": self.image_data
        }