src/models.py

"""
Data Models for Character Sheet Pro
====================================

Dataclasses for generation requests and results.
"""

from dataclasses import dataclass, field
from typing import Optional, List, Dict, Any
from PIL import Image
from datetime import datetime


@dataclass
class GenerationRequest:
    """Request for image generation."""

    prompt: str
    input_images: List[Image.Image] = field(default_factory=list)
    aspect_ratio: str = "1:1"
    temperature: float = 0.4
    negative_prompt: Optional[str] = None

    @property
    def has_input_images(self) -> bool:
        """Check if request has input images."""
        return len(self.input_images) > 0


@dataclass
class GenerationResult:
    """Result from image generation."""

    success: bool
    image: Optional[Image.Image] = None
    message: str = ""
    generation_time: Optional[float] = None

    @classmethod
    def success_result(
        cls,
        image: Image.Image,
        message: str = "Generated successfully",
        generation_time: Optional[float] = None
    ) -> "GenerationResult":
        """Create successful result."""
        return cls(
            success=True,
            image=image,
            message=message,
            generation_time=generation_time
        )

    @classmethod
    def error_result(cls, message: str) -> "GenerationResult":
        """Create error result."""
        return cls(success=False, message=message)


@dataclass
class CharacterSheetConfig:
    """Configuration for character sheet generation."""

    # Aspect ratios
    face_aspect_ratio: str = "3:4"  # 864x1184
    body_aspect_ratio: str = "9:16"  # 768x1344

    # Generation temperatures
    face_temperature: float = 0.35
    body_temperature: float = 0.35
    normalize_temperature: float = 0.5

    # Layout
    spacing: int = 20
    background_color: str = "#2C2C2C"

    # Retry settings
    max_retries: int = 3
    retry_delay: float = 30.0
    rate_limit_delay_min: float = 2.0
    rate_limit_delay_max: float = 3.0


@dataclass
class CharacterSheetMetadata:
    """Metadata for generated character sheet."""

    character_name: str
    input_type: str  # "Face Only", "Full Body", "Face + Body"
    costume_description: str
    backend: str
    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())
    views: int = 7
    stages: Dict[str, Any] = field(default_factory=dict)