imageforge/backend/app/api/schemas.py

from __future__ import annotations

from datetime import datetime
from enum import Enum
from typing import Literal

from pydantic import BaseModel, Field, field_validator


class JobStatus(str, Enum):
    QUEUED = "queued"
    RUNNING = "running"
    DONE = "done"
    ERROR = "error"
    CANCELLED = "cancelled"


class GenerateRequest(BaseModel):
    prompt: str = Field(min_length=1, max_length=2000)
    negative_prompt: str = Field(default="", max_length=2000)
    model: str = Field(default="localai")  # Changed from "dummy" to "localai"
    size: Literal["512x512", "768x768", "1024x1024", "1024x1536", "1536x1024"] = "512x512"
    count: int = Field(default=1, ge=1, le=4)
    seed: int | None = Field(default=None, ge=0)
    random_seed: bool = True
    steps: int = Field(default=30, ge=1, le=100)
    guidance: float = Field(default=6.5, ge=1.0, le=20.0)
    image_type: Literal[
        "general",
        "photo",
        "portrait",
        "landscape",
        "architecture",
        "product",
        "logo",
        "icon",
        "poster",
        "illustration",
        "anime",
        "pixel_art",
        "sketch",
        "painting",
        "3d",
    ] = "general"
    style_preset: Literal[
        "auto",
        "photorealistic",
        "cinematic",
        "minimal",
        "vibrant",
        "monochrome",
        "watercolor",
        "oil",
        "noir",
        "fantasy",
    ] = "auto"
    style_strength: int = Field(default=60, ge=0, le=100)
    admin_override: bool = False
    init_image_path: str | None = None
    img2img_strength: float = Field(default=0.45, ge=0.0, le=1.0)
    model_variant: str | None = Field(default=None, max_length=200)

    @field_validator("prompt")
    @classmethod
    def validate_prompt(cls, value: str) -> str:
        value = value.strip()
        if not value:
            raise ValueError("Prompt must not be empty")
        return value


class GenerateResponse(BaseModel):
    job_id: str


class RetryResponse(BaseModel):
    old_job_id: str
    new_job_id: str


class JobInfoResponse(BaseModel):
    job_id: str
    status: JobStatus
    progress: int
    message: str
    created_at: datetime
    updated_at: datetime
    image_paths: list[str] = Field(default_factory=list)
    output_dir: str | None = None
    error: str | None = None


class CancelResponse(BaseModel):
    success: bool


class ModelInfo(BaseModel):
    id: str
    name: str
    available: bool
    description: str


class HistoryItem(BaseModel):
    prompt: str
    negative_prompt: str
    timestamp: datetime


class HealthResponse(BaseModel):
    status: str
    timestamp: datetime


class PresetPayload(BaseModel):
    name: str = Field(min_length=1, max_length=80)
    prompt: str = ""
    negative_prompt: str = ""
    model: str = "dummy"
    size: Literal["512x512", "768x768", "1024x1024", "1024x1536", "1536x1024"] = "512x512"
    count: int = Field(default=1, ge=1, le=4)
    steps: int = Field(default=30, ge=1, le=100)
    guidance: float = Field(default=6.5, ge=1.0, le=20.0)
    image_type: str = "general"
    style_preset: str = "auto"
    style_strength: int = Field(default=60, ge=0, le=100)


class PresetResponse(PresetPayload):
    updated_at: datetime


class DashboardStats(BaseModel):
    queued: int
    running: int
    done: int
    error: int
    cancelled: int
    total: int
    last_24h: int


class ExportRequest(BaseModel):
    source_path: str
    format: Literal["png", "jpg", "webp"] = "png"
    quality: int = Field(default=92, ge=1, le=100)
    max_width: int | None = Field(default=None, ge=64, le=8192)
    max_height: int | None = Field(default=None, ge=64, le=8192)


class ExportResponse(BaseModel):
    output_path: str


class AdminSettings(BaseModel):
    content_profile: str
    rate_limit_per_minute: int
    output_retention_days: int
    adult_enabled: bool = False


class MetricsResponse(BaseModel):
    metrics: dict[str, float | int]