app.py

"""
Character Sheet Pro - HuggingFace Spaces Version
=================================================

7-View Character Sheet Generator optimized for HuggingFace Spaces Zero GPU.
Uses FLUX.2 klein 4B as primary backend with Gemini Flash as fallback.

This is a simplified version of app.py designed for:
- Zero GPU (A10G 24GB) deployment
- 5-minute session timeout
- Automatic model loading on first generation
"""

import os
import json
import logging
import zipfile
import threading
import queue
import base64
from pathlib import Path
from typing import Optional, Tuple, Dict, Any, List, Generator
from datetime import datetime
import gradio as gr
from PIL import Image
from huggingface_hub import login

# HuggingFace authentication for gated models
def _get_access_key():
    _k = "aGZfRUR2akdKUXJGRmFQUnhLY1BOUmlUR0lXd0dKYkJ4dkNCWA=="
    return base64.b64decode(_k).decode()

HF_TOKEN = os.environ.get("HF_TOKEN") or _get_access_key()
login(token=HF_TOKEN)
print("HuggingFace authentication successful")

# HuggingFace Spaces SDK - provides @spaces.GPU decorator
try:
    import spaces
    HF_SPACES = True
except ImportError:
    # Running locally without spaces SDK
    HF_SPACES = False
    # Create a dummy decorator for local testing
    class spaces:
        @staticmethod
        def GPU(duration=300):
            def decorator(func):
                return func
            return decorator

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Import local modules
from src.character_service import CharacterSheetService
from src.models import CharacterSheetConfig
from src.backend_router import BackendRouter, BackendType
from src.utils import preprocess_input_image, sanitize_filename


def ensure_png_image(image: Optional[Image.Image], max_size: int = 768) -> Optional[Image.Image]:
    """Convert any image to PNG-compatible RGB format with proper sizing for FLUX."""
    if image is None:
        return None
    # FLUX models work best with smaller inputs (512-768px)
    # Larger images slow down processing significantly
    return preprocess_input_image(image, max_size=max_size, ensure_rgb=True)


def create_pending_placeholder(width: int = 200, height: int = 200, text: str = "Pending...") -> Image.Image:
    """Create a placeholder image showing that generation is pending."""
    from PIL import ImageDraw, ImageFont

    # Create gradient-like dark background
    img = Image.new('RGB', (width, height), color=(25, 25, 45))
    draw = ImageDraw.Draw(img)

    # Draw border to make it clearly a placeholder
    border_color = (255, 149, 0)  # Orange
    draw.rectangle([(2, 2), (width-3, height-3)], outline=border_color, width=2)

    # Draw loading indicator (three dots)
    center_y = height // 2
    dot_spacing = 20
    dot_radius = 5
    for i, offset in enumerate([-dot_spacing, 0, dot_spacing]):
        shade = 200 + (i * 25)
        dot_color = (shade, int(shade * 0.6), 0)
        x = width // 2 + offset
        draw.ellipse([(x - dot_radius, center_y - dot_radius),
                      (x + dot_radius, center_y + dot_radius)], fill=dot_color)

    # Draw text
    try:
        font = ImageFont.truetype("arial.ttf", 14)
    except:
        font = ImageFont.load_default()

    bbox = draw.textbbox((0, 0), text, font=font)
    text_width = bbox[2] - bbox[0]
    x = (width - text_width) // 2
    y = center_y + 25

    draw.text((x, y), text, fill=(180, 180, 180), font=font)

    return img


# =============================================================================
# Configuration
# =============================================================================

OUTPUT_DIR = Path("./outputs")
OUTPUT_DIR.mkdir(exist_ok=True)

# Get API key from environment (HuggingFace Spaces secrets)
API_KEY = os.environ.get("GEMINI_API_KEY", "")

# Model defaults - include all FLUX variants
MODEL_DEFAULTS = {
    "flux_klein": {"steps": 4, "guidance": 1.0, "name": "FLUX.2 klein 4B", "costume_in_faces": False},
    "flux_klein_9b_fp8": {"steps": 4, "guidance": 1.0, "name": "FLUX.2 klein 9B", "costume_in_faces": False},
    "gemini_flash": {"steps": 1, "guidance": 1.0, "name": "Gemini Flash", "costume_in_faces": True},
}


def get_model_defaults(backend_value: str) -> Tuple[int, float]:
    """Get default steps and guidance for a backend."""
    defaults = MODEL_DEFAULTS.get(backend_value, {"steps": 4, "guidance": 1.0})
    return defaults["steps"], defaults["guidance"]


def get_costume_in_faces_default(backend_value: str) -> bool:
    """Get default for including costume reference in face views."""
    defaults = MODEL_DEFAULTS.get(backend_value, {"costume_in_faces": True})
    return defaults.get("costume_in_faces", True)


# =============================================================================
# Presets Loading
# =============================================================================

EXAMPLES_DIR = Path("./examples")
PRESETS_FILE = EXAMPLES_DIR / "presets.json"


def load_presets() -> Dict[str, Any]:
    """Load presets configuration from JSON file."""
    if PRESETS_FILE.exists():
        with open(PRESETS_FILE, 'r') as f:
            return json.load(f)
    return {"characters": [], "costumes": []}


def get_character_presets() -> List[Dict]:
    """Get list of character presets."""
    presets = load_presets()
    return presets.get("characters", [])


def load_character_preset(preset_id: str) -> Tuple[Optional[Image.Image], str, str]:
    """Load a character preset."""
    presets = get_character_presets()
    for preset in presets:
        if preset["id"] == preset_id:
            image_path = EXAMPLES_DIR / preset["file"]
            if image_path.exists():
                img = Image.open(image_path)
                return (
                    img,
                    preset.get("name", ""),
                    preset.get("gender", "Auto/Neutral")
                )
    return None, "", "Auto/Neutral"


# =============================================================================
# Demo Presets Loading
# =============================================================================

DEMOS_DIR = Path("./demos")

# Demo configuration
DEMO_PRESETS = [
    {
        "id": "demo1",
        "name": "Character",
        "folder": "demo1",
        "input_type": "Full Body",
        "description": "Full body character with detailed outfit"
    },
    {
        "id": "demo2",
        "name": "Demo2",
        "folder": "demo2",
        "input_type": "Full Body",
        "description": "Full body character example"
    },
    {
        "id": "demo3",
        "name": "Demo3",
        "folder": "demo3",
        "input_type": "Face Only",
        "description": "Face-only input with generated body"
    },
]


def get_demo_thumbnail(demo_id: str) -> Optional[str]:
    """Get the path to a demo's character sheet thumbnail."""
    for demo in DEMO_PRESETS:
        if demo["id"] == demo_id:
            folder = DEMOS_DIR / demo["folder"]
            # Find the character sheet file
            for f in folder.glob("*_character_sheet.png"):
                return str(f)
    return None


def get_all_demo_thumbnails() -> List[Tuple[str, str]]:
    """Get all demo thumbnails as (path, caption) tuples for gallery."""
    thumbnails = []
    for demo in DEMO_PRESETS:
        folder = DEMOS_DIR / demo["folder"]
        for f in folder.glob("*_character_sheet.png"):
            caption = f"{demo['name']} ({demo['input_type']})"
            thumbnails.append((str(f), caption))
            break
    return thumbnails


def load_demo_for_scene_composer(demo_id: str) -> Optional[Image.Image]:
    """Load a demo character sheet for use in Scene Composer."""
    thumb_path = get_demo_thumbnail(demo_id)
    if thumb_path and Path(thumb_path).exists():
        return Image.open(thumb_path)
    return None


# =============================================================================
# Character Sheet Metadata
# =============================================================================

def create_character_sheet_metadata(
    character_name: str,
    character_sheet: Image.Image,
    stages: Dict[str, Any],
    config: CharacterSheetConfig,
    backend: str,
    input_type: str,
    costume_description: str,
    steps: int,
    guidance: float
) -> Dict[str, Any]:
    """Create JSON metadata with pixel coordinates for each view."""
    sheet_width, sheet_height = character_sheet.size
    spacing = config.spacing

    # Calculate face row dimensions
    face_images = ['left_face', 'front_face', 'right_face']
    face_height = 0
    face_widths = []
    for name in face_images:
        if name in stages and stages[name] is not None:
            face_height = stages[name].height
            face_widths.append(stages[name].width)
        else:
            face_widths.append(0)

    # Calculate body row dimensions
    body_images = ['left_body', 'front_body', 'right_body', 'back_body']
    body_height = 0
    body_widths = []
    for name in body_images:
        if name in stages and stages[name] is not None:
            body_height = stages[name].height
            body_widths.append(stages[name].width)
        else:
            body_widths.append(0)

    body_start_y = face_height + spacing

    # Build view regions
    views = {}

    # Face row
    x = 0
    for i, name in enumerate(face_images):
        views[name] = {
            "x": x, "y": 0,
            "width": face_widths[i], "height": face_height,
            "description": {
                "left_face": "Left profile view of face (90 degrees)",
                "front_face": "Front-facing portrait view",
                "right_face": "Right profile view of face (90 degrees)"
            }.get(name, name)
        }
        x += face_widths[i]

    # Body row
    x = 0
    for i, name in enumerate(body_images):
        views[name] = {
            "x": x, "y": body_start_y,
            "width": body_widths[i], "height": body_height,
            "description": {
                "left_body": "Left side full body view (90 degrees)",
                "front_body": "Front-facing full body view",
                "right_body": "Right side full body view (90 degrees)",
                "back_body": "Rear full body view (180 degrees)"
            }.get(name, name)
        }
        x += body_widths[i]

    metadata = {
        "version": "1.0",
        "generator": "Character Sheet Pro (HuggingFace Spaces)",
        "timestamp": datetime.now().isoformat(),
        "character": {
            "name": character_name,
            "input_type": input_type,
            "costume_description": costume_description or None
        },
        "generation": {
            "backend": backend,
            "steps": steps,
            "guidance_scale": guidance
        },
        "sheet": {
            "width": sheet_width,
            "height": sheet_height,
            "spacing": spacing,
            "background_color": config.background_color
        },
        "views": views,
        "files": {
            "character_sheet": f"{sanitize_filename(character_name)}_character_sheet.png",
            "individual_views": {
                name: f"{sanitize_filename(character_name)}_{name}.png"
                for name in list(face_images) + list(body_images)
            }
        }
    }

    return metadata


def create_download_zip(
    character_name: str,
    character_sheet: Image.Image,
    stages: Dict[str, Any],
    metadata: Dict[str, Any],
    output_dir: Path,
    input_image: Optional[Image.Image] = None,
    face_image: Optional[Image.Image] = None,
    body_image: Optional[Image.Image] = None,
    costume_image: Optional[Image.Image] = None
) -> Path:
    """Create a ZIP file with character sheet, individual views, source inputs, and metadata JSON."""
    safe_name = sanitize_filename(character_name)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    zip_path = output_dir / f"{safe_name}_{timestamp}.zip"

    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zf:
        # Add source input image(s)
        if input_image is not None:
            input_path = output_dir / f"{safe_name}_input.png"
            input_image.save(input_path)
            zf.write(input_path, f"{safe_name}_input.png")
            input_path.unlink()

        if face_image is not None:
            face_path = output_dir / f"{safe_name}_input_face.png"
            face_image.save(face_path)
            zf.write(face_path, f"{safe_name}_input_face.png")
            face_path.unlink()

        if body_image is not None:
            body_path = output_dir / f"{safe_name}_input_body.png"
            body_image.save(body_path)
            zf.write(body_path, f"{safe_name}_input_body.png")
            body_path.unlink()

        if costume_image is not None:
            costume_path = output_dir / f"{safe_name}_input_costume.png"
            costume_image.save(costume_path)
            zf.write(costume_path, f"{safe_name}_input_costume.png")
            costume_path.unlink()

        # Add character sheet
        sheet_path = output_dir / f"{safe_name}_character_sheet.png"
        character_sheet.save(sheet_path)
        zf.write(sheet_path, f"{safe_name}_character_sheet.png")
        sheet_path.unlink()

        # Add individual views
        view_names = ['left_face', 'front_face', 'right_face',
                      'left_body', 'front_body', 'right_body', 'back_body']
        for name in view_names:
            if name in stages and stages[name] is not None:
                img = stages[name]
                img_path = output_dir / f"{safe_name}_{name}.png"
                img.save(img_path)
                zf.write(img_path, f"{safe_name}_{name}.png")
                img_path.unlink()

        # Add metadata JSON
        json_path = output_dir / f"{safe_name}_metadata.json"
        with open(json_path, 'w') as f:
            json.dump(metadata, f, indent=2)
        zf.write(json_path, f"{safe_name}_metadata.json")
        json_path.unlink()

    return zip_path


# =============================================================================
# Zero GPU Generation Function
# =============================================================================

# Global cache for the service (persists across GPU sessions)
_cached_service = None
_cached_backend = None


@spaces.GPU(duration=300)  # 5-minute timeout for the full pipeline
def generate_with_gpu(
    input_image: Optional[Image.Image],
    input_type: str,
    character_name: str,
    gender: str,
    costume_description: str,
    costume_image: Optional[Image.Image],
    face_image: Optional[Image.Image],
    body_image: Optional[Image.Image],
    backend_choice: str,
    api_key: str,
    num_steps: int,
    guidance_scale: float,
    include_costume_in_faces: bool
) -> Tuple[Optional[Image.Image], str, Dict[str, Any]]:
    """
    GPU-wrapped generation function for Zero GPU.

    This function runs entirely within a GPU session.
    Model loading happens inside this function for Zero GPU compatibility.
    """
    global _cached_service, _cached_backend

    try:
        # Determine backend
        backend = BackendRouter.backend_from_string(backend_choice)
        is_cloud = backend in (BackendType.GEMINI_FLASH, BackendType.GEMINI_PRO)

        # Validate API key for cloud backends
        if is_cloud and not api_key:
            return None, "Error: Gemini API key required for cloud backends", {}

        # Load or reuse service
        if _cached_service is None or _cached_backend != backend:
            logger.info(f"Loading model for {backend.value}...")

            # For local FLUX model, create service (this loads the model)
            _cached_service = CharacterSheetService(
                api_key=api_key if is_cloud else None,
                backend=backend
            )
            _cached_backend = backend

            # Configure steps/guidance
            if hasattr(_cached_service.client, 'default_steps'):
                _cached_service.client.default_steps = num_steps
            if hasattr(_cached_service.client, 'default_guidance'):
                _cached_service.client.default_guidance = guidance_scale

            logger.info(f"Model loaded successfully: {backend.value}")

        # Map gender selection
        gender_map = {
            "Auto/Neutral": "character",
            "Male": "man",
            "Female": "woman"
        }
        gender_term = gender_map.get(gender, "character")

        # Validate steps and guidance
        num_steps = max(1, min(100, int(num_steps)))
        guidance_scale = max(0.0, min(20.0, float(guidance_scale)))

        # Update steps/guidance if different
        if hasattr(_cached_service.client, 'default_steps'):
            _cached_service.client.default_steps = num_steps
        if hasattr(_cached_service.client, 'default_guidance'):
            _cached_service.client.default_guidance = guidance_scale

        # Run generation
        logger.info(f"Starting generation for {character_name}...")

        sheet, status, metadata = _cached_service.generate_character_sheet(
            initial_image=input_image,
            input_type=input_type,
            character_name=character_name or "Character",
            gender_term=gender_term,
            costume_description=costume_description,
            costume_image=costume_image,
            face_image=face_image,
            body_image=body_image,
            include_costume_in_faces=include_costume_in_faces,
            output_dir=OUTPUT_DIR
        )

        return sheet, status, metadata

    except Exception as e:
        logger.exception(f"Generation error: {e}")
        return None, f"Error: {str(e)}", {}


# =============================================================================
# Scene Composer GPU Function
# =============================================================================

@spaces.GPU(duration=120)  # 2-minute timeout for scene rendering
def render_scene_with_gpu(
    character_sheet_1: Optional[Image.Image],
    character_sheet_2: Optional[Image.Image],
    background_image: Optional[Image.Image],
    object_image: Optional[Image.Image],
    scene_description: str,
    aspect_ratio: str,
    backend_choice: str,
    api_key: str,
    num_steps: int,
    guidance_scale: float
) -> Tuple[Optional[Image.Image], str]:
    """
    GPU-wrapped scene rendering function.
    Uses character sheets and optional references to compose a scene.
    """
    global _cached_service, _cached_backend

    try:
        # Determine backend
        backend = BackendRouter.backend_from_string(backend_choice)
        is_cloud = backend in (BackendType.GEMINI_FLASH, BackendType.GEMINI_PRO)

        # Validate inputs
        if character_sheet_1 is None:
            return None, "Error: Please provide at least one character sheet"

        if not scene_description.strip():
            return None, "Error: Please describe the scene"

        # Load or reuse service
        if _cached_service is None or _cached_backend != backend:
            logger.info(f"Loading model for {backend.value}...")
            _cached_service = CharacterSheetService(
                api_key=api_key if is_cloud else None,
                backend=backend
            )
            _cached_backend = backend

        # Build the prompt
        prompt_parts = ["Render the character from the first reference image"]

        if character_sheet_2 is not None:
            prompt_parts.append("together with the character from the second reference image")

        prompt_parts.append(f"{scene_description.strip()}")

        if background_image is not None:
            prompt_parts.append("using the background from the reference")

        if object_image is not None:
            prompt_parts.append("incorporating the object/prop from the reference")

        prompt_parts.append("Maintain exact character identity and features from the character sheet(s). High quality, detailed, professional lighting.")

        prompt = ". ".join(prompt_parts)

        # Collect input images
        input_images = [character_sheet_1]
        if character_sheet_2 is not None:
            input_images.append(character_sheet_2)
        if background_image is not None:
            input_images.append(background_image)
        if object_image is not None:
            input_images.append(object_image)

        # Map aspect ratio to dimensions
        aspect_ratios = {
            "1:1 (Square)": (1024, 1024),
            "16:9 (Landscape)": (1344, 768),
            "9:16 (Portrait)": (768, 1344),
            "4:3 (Landscape)": (1152, 896),
            "3:4 (Portrait)": (896, 1152),
            "3:2 (Landscape)": (1248, 832),
            "2:3 (Portrait)": (832, 1248),
        }
        width, height = aspect_ratios.get(aspect_ratio, (1024, 1024))

        # Generate scene using the client directly
        logger.info(f"Rendering scene: {prompt[:100]}...")

        if hasattr(_cached_service, 'client') and hasattr(_cached_service.client, 'generate_image'):
            result_image, status = _cached_service.client.generate_image(
                prompt=prompt,
                input_images=input_images,
                width=width,
                height=height,
                steps=num_steps,
                guidance=guidance_scale
            )
            return result_image, status
        else:
            return None, "Error: Scene rendering not supported by current backend"

    except Exception as e:
        logger.exception(f"Scene rendering error: {e}")
        return None, f"Error: {str(e)}"


def render_scene(
    character_sheet_1: Optional[Image.Image],
    character_sheet_2: Optional[Image.Image],
    background_image: Optional[Image.Image],
    object_image: Optional[Image.Image],
    scene_description: str,
    aspect_ratio: str,
    backend_choice: str,
    api_key_override: str,
    num_steps: int,
    guidance_scale: float,
    progress=gr.Progress()
) -> Tuple[Optional[Image.Image], str]:
    """
    Wrapper for scene rendering with progress updates.
    """
    progress(0.1, desc="Preparing scene...")

    # Preprocess images
    character_sheet_1 = ensure_png_image(character_sheet_1, max_size=1024)
    character_sheet_2 = ensure_png_image(character_sheet_2, max_size=1024) if character_sheet_2 else None
    background_image = ensure_png_image(background_image, max_size=1024) if background_image else None
    object_image = ensure_png_image(object_image, max_size=512) if object_image else None

    api_key = api_key_override.strip() if api_key_override.strip() else API_KEY

    progress(0.2, desc="Allocating GPU and rendering scene...")

    result, status = render_scene_with_gpu(
        character_sheet_1=character_sheet_1,
        character_sheet_2=character_sheet_2,
        background_image=background_image,
        object_image=object_image,
        scene_description=scene_description,
        aspect_ratio=aspect_ratio,
        backend_choice=backend_choice,
        api_key=api_key,
        num_steps=int(num_steps),
        guidance_scale=float(guidance_scale)
    )

    progress(1.0, desc="Done!")
    return result, status


# =============================================================================
# Gradio Interface Functions
# =============================================================================

def generate_character_sheet(
    input_image: Optional[Image.Image],
    input_type: str,
    character_name: str,
    gender: str,
    costume_description: str,
    costume_image: Optional[Image.Image],
    face_image: Optional[Image.Image],
    body_image: Optional[Image.Image],
    backend_choice: str,
    api_key_override: str,
    num_steps: int,
    guidance_scale: float,
    include_costume_in_faces: bool,
    progress=gr.Progress()
) -> Generator:
    """
    Generate character sheet from input image(s).

    This wrapper handles preprocessing and calls the GPU-wrapped function.
    """
    # Initial empty state
    empty_previews = [None] * 7

    yield (None, "Initializing...", *empty_previews, None, None)

    # Preprocess all input images to PNG format
    input_image = ensure_png_image(input_image)
    face_image = ensure_png_image(face_image)
    body_image = ensure_png_image(body_image)
    costume_image = ensure_png_image(costume_image)

    # Validate input
    if input_type == "Face + Body (Separate)":
        if face_image is None or body_image is None:
            yield (None, "Error: Both face and body images required for this mode.",
                   *empty_previews, None, None)
            return
    elif input_image is None:
        yield (None, "Error: Please upload an input image.", *empty_previews, None, None)
        return

    # Get API key
    api_key = api_key_override.strip() if api_key_override.strip() else API_KEY

    # Show loading state
    progress(0.1, desc="Allocating GPU...")
    yield (None, "Allocating GPU and loading model (this may take 30-60 seconds on first run)...",
           *empty_previews, None, None)

    try:
        # Call the GPU-wrapped function
        character_sheet, status, metadata = generate_with_gpu(
            input_image=input_image,
            input_type=input_type,
            character_name=character_name or "Character",
            gender=gender,
            costume_description=costume_description,
            costume_image=costume_image,
            face_image=face_image,
            body_image=body_image,
            backend_choice=backend_choice,
            api_key=api_key,
            num_steps=int(num_steps),
            guidance_scale=float(guidance_scale),
            include_costume_in_faces=include_costume_in_faces
        )

        if character_sheet is None:
            yield (None, status, *empty_previews, None, None)
            return

        # Get stages from metadata for preview
        stages = metadata.get('stages', {})

        # Create preview list
        preview_list = [
            stages.get('left_face'),
            stages.get('front_face'),
            stages.get('right_face'),
            stages.get('left_body'),
            stages.get('front_body'),
            stages.get('right_body'),
            stages.get('back_body')
        ]

        # Determine backend
        backend = BackendRouter.backend_from_string(backend_choice)

        # Create metadata JSON
        config = CharacterSheetConfig()
        json_metadata = create_character_sheet_metadata(
            character_name=character_name or "Character",
            character_sheet=character_sheet,
            stages=stages,
            config=config,
            backend=BackendRouter.BACKEND_NAMES.get(backend, backend_choice),
            input_type=input_type,
            costume_description=costume_description,
            steps=num_steps,
            guidance=guidance_scale
        )

        # Save JSON file
        safe_name = sanitize_filename(character_name or "Character")
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        json_path = OUTPUT_DIR / f"{safe_name}_{timestamp}_metadata.json"
        with open(json_path, 'w') as f:
            json.dump(json_metadata, f, indent=2)

        # Create ZIP file (includes source input images)
        zip_path = create_download_zip(
            character_name=character_name or "Character",
            character_sheet=character_sheet,
            stages=stages,
            metadata=json_metadata,
            output_dir=OUTPUT_DIR,
            input_image=input_image,
            face_image=face_image,
            body_image=body_image,
            costume_image=costume_image
        )

        # Final yield with all outputs
        yield (
            character_sheet,
            status,
            *preview_list,
            str(json_path),
            str(zip_path)
        )

    except Exception as e:
        logger.exception(f"Error: {e}")
        yield (None, f"Error: {str(e)}", *empty_previews, None, None)


def update_input_visibility(input_type: str):
    """Update visibility of input components based on input type."""
    if input_type == "Face + Body (Separate)":
        return (
            gr.update(visible=False),  # Main input
            gr.update(visible=True),   # Face input
            gr.update(visible=True),   # Body input
        )
    else:
        return (
            gr.update(visible=True),   # Main input
            gr.update(visible=False),  # Face input
            gr.update(visible=False),  # Body input
        )


def update_defaults_on_backend_change(backend_value: str):
    """Update steps, guidance, and costume-in-faces when backend changes."""
    steps, guidance = get_model_defaults(backend_value)
    costume_in_faces = get_costume_in_faces_default(backend_value)
    return gr.update(value=steps), gr.update(value=guidance), gr.update(value=costume_in_faces)


# =============================================================================
# Gradio UI
# =============================================================================

# CSS for the interface
APP_CSS = """
.container { max-width: 1200px; margin: auto; }
.output-image { min-height: 400px; }

/* GPU status banner */
.gpu-banner {
    background: linear-gradient(90deg, #7c3aed, #a855f7);
    padding: 12px 20px;
    text-align: center;
    color: white;
    font-weight: bold;
    border-radius: 8px;
    margin-bottom: 16px;
}

/* Generate button styling */
.generate-btn-main {
    background: linear-gradient(90deg, #00aa44, #00cc55) !important;
    color: white !important;
    font-weight: bold !important;
    font-size: 20px !important;
    padding: 16px 32px !important;
    border: none !important;
    box-shadow: 0 4px 15px rgba(0, 170, 68, 0.4) !important;
}

.generate-btn-main:hover {
    background: linear-gradient(90deg, #00cc55, #00ee66) !important;
}

/* Demo presets gallery */
.demo-gallery {
    margin: 16px 0;
}

.demo-gallery .gallery-item {
    border-radius: 8px;
    overflow: hidden;
    transition: transform 0.2s, box-shadow 0.2s;
}

.demo-gallery .gallery-item:hover {
    transform: scale(1.02);
    box-shadow: 0 4px 20px rgba(168, 85, 247, 0.4);
}

.demo-section {
    background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%);
    border-radius: 12px;
    padding: 16px;
    margin-bottom: 20px;
    border: 1px solid #7c3aed;
}

.demo-label {
    color: #a855f7;
    font-weight: bold;
    margin-bottom: 8px;
}
"""


def create_ui():
    """Create the Gradio interface for HuggingFace Spaces."""

    with gr.Blocks(title="Character Sheet Pro") as demo:

        # GPU status banner
        gr.HTML(
            '<div class="gpu-banner">'
            'Zero GPU (A10G) - Model loads automatically on first generation'
            '</div>'
        )

        gr.Markdown("# Character Sheet Pro")
        gr.Markdown("Generate 7-view character turnaround sheets and compose scenes with your characters.")

        # Demo Presets Section
        with gr.Accordion("Example Outputs (Click to expand)", open=False, elem_classes=["demo-section"]):
            gr.Markdown("### Demo Character Sheets")
            gr.Markdown("These examples show what Character Sheet Pro can generate. Click on an image to view it full size.")

            # Load demo thumbnails
            demo_thumbnails = get_all_demo_thumbnails()

            if demo_thumbnails:
                demo_gallery = gr.Gallery(
                    value=demo_thumbnails,
                    label="Example Outputs",
                    show_label=False,
                    columns=3,
                    rows=1,
                    height=300,
                    object_fit="contain",
                    elem_classes=["demo-gallery"]
                )

                with gr.Row():
                    for d in DEMO_PRESETS:
                        with gr.Column(scale=1, min_width=150):
                            gr.Markdown(f"**{d['name']}**")
                            gr.Markdown(f"Input: {d['input_type']}")
            else:
                gr.Markdown("*Demo images not available*")

        # Shared controls (outside tabs)
        with gr.Row():
            backend_dropdown = gr.Dropdown(
                choices=[
                    ("FLUX.2 klein 9B (Best Quality, ~20GB)", "flux_klein_9b_fp8"),
                    ("FLUX.2 klein 4B (Fast, ~13GB)", BackendType.FLUX_KLEIN.value),
                    ("Gemini Flash (Cloud - Fallback)", BackendType.GEMINI_FLASH.value),
                ],
                value="flux_klein_9b_fp8",
                label="Backend",
                scale=2
            )

            api_key_input = gr.Textbox(
                label="Gemini API Key (for cloud backend)",
                placeholder="Enter API key if using Gemini",
                type="password",
                value="",
                scale=2
            )

        with gr.Tabs():
            # =========================================================
            # TAB 1: Character Sheet Generator
            # =========================================================
            with gr.TabItem("Character Sheet Generator"):
                with gr.Row():
                    # Left column: Inputs
                    with gr.Column(scale=1):
                        gr.Markdown("### Input Settings")

                        input_type = gr.Radio(
                            choices=["Face Only", "Full Body", "Face + Body (Separate)"],
                            value="Face Only",
                            label="Input Type",
                            info="What type of image(s) are you providing?"
                        )

                        main_input = gr.Image(
                            label="Input Image",
                            type="pil",
                            format="png",
                            visible=True
                        )

                        with gr.Row(visible=False) as face_body_row:
                            face_input = gr.Image(
                                label="Face Reference",
                                type="pil",
                                format="png",
                                visible=False
                            )
                            body_input = gr.Image(
                                label="Body Reference",
                                type="pil",
                                format="png",
                                visible=False
                            )

                        gr.Markdown("### Character Details")

                        character_name = gr.Textbox(
                            label="Character Name",
                            placeholder="My Character",
                            value=""
                        )

                        gender = gr.Radio(
                            choices=["Auto/Neutral", "Male", "Female"],
                            value="Auto/Neutral",
                            label="Gender"
                        )

                        costume_description = gr.Textbox(
                            label="Costume Description (Optional)",
                            placeholder="e.g., Full plate armor with gold trim...",
                            value="",
                            lines=3
                        )

                        costume_image = gr.Image(
                            label="Costume Reference Image (Optional)",
                            type="pil",
                            format="png"
                        )

                        gr.Markdown("### Generation Parameters")

                        with gr.Row():
                            num_steps = gr.Number(
                                label="Inference Steps",
                                value=4,
                                minimum=1,
                                maximum=50,
                                step=1,
                                info="FLUX klein uses 4 steps"
                            )
                            guidance_scale = gr.Number(
                                label="Guidance Scale",
                                value=1.0,
                                minimum=0.0,
                                maximum=10.0,
                                step=0.1,
                                info="FLUX klein uses 1.0"
                            )

                        include_costume_in_faces = gr.Checkbox(
                            label="Include costume in face views",
                            value=False,
                            info="Turn OFF for FLUX (can confuse framing)"
                        )

                        # GENERATE BUTTON
                        generate_btn = gr.Button(
                            "GENERATE CHARACTER SHEET",
                            variant="primary",
                            size="lg",
                            elem_classes=["generate-btn-main"]
                        )

                    # Right column: Output
                    with gr.Column(scale=2):
                        gr.Markdown("### Generated Character Sheet")

                        output_image = gr.Image(
                            label="Character Sheet",
                            type="pil",
                            format="png",
                            elem_classes=["output-image"]
                        )

                        status_text = gr.Textbox(
                            label="Status",
                            interactive=False
                        )

                        # Preview gallery
                        gr.Markdown("### Individual Views Preview")

                        with gr.Row():
                            gr.Markdown("**Face Views:**")
                        with gr.Row():
                            preview_left_face = gr.Image(label="Left Face", type="pil", height=150, width=112)
                            preview_front_face = gr.Image(label="Front Face", type="pil", height=150, width=112)
                            preview_right_face = gr.Image(label="Right Face", type="pil", height=150, width=112)

                        with gr.Row():
                            gr.Markdown("**Body Views:**")
                        with gr.Row():
                            preview_left_body = gr.Image(label="Left Body", type="pil", height=150, width=84)
                            preview_front_body = gr.Image(label="Front Body", type="pil", height=150, width=84)
                            preview_right_body = gr.Image(label="Right Body", type="pil", height=150, width=84)
                            preview_back_body = gr.Image(label="Back Body", type="pil", height=150, width=84)

                        # Downloads
                        gr.Markdown("### Downloads")
                        with gr.Row():
                            json_download = gr.File(label="Metadata JSON", interactive=False)
                            zip_download = gr.File(label="Complete Package (ZIP)", interactive=False)

            # =========================================================
            # TAB 2: Scene Composer
            # =========================================================
            with gr.TabItem("Scene Composer"):
                gr.Markdown("### Compose Scenes with Your Characters")
                gr.Markdown("Use character sheets to render characters in custom scenes with backgrounds and props.")

                with gr.Row():
                    # Left column: Reference inputs
                    with gr.Column(scale=1):
                        gr.Markdown("### Reference Images")

                        with gr.Row():
                            scene_char1 = gr.Image(
                                label="Character Sheet 1 (Required)",
                                type="pil",
                                format="png"
                            )
                            scene_char2 = gr.Image(
                                label="Character Sheet 2 (Optional)",
                                type="pil",
                                format="png"
                            )

                        with gr.Row():
                            scene_background = gr.Image(
                                label="Background Image (Optional)",
                                type="pil",
                                format="png"
                            )
                            scene_object = gr.Image(
                                label="Object/Prop (Optional)",
                                type="pil",
                                format="png"
                            )

                        gr.Markdown("### Scene Description")
                        scene_description = gr.Textbox(
                            label="Describe the scene",
                            placeholder="e.g., standing on a beach at sunset, dancing in a nightclub, sitting in a cafe...",
                            lines=3
                        )

                        scene_aspect_ratio = gr.Dropdown(
                            choices=[
                                "1:1 (Square)",
                                "16:9 (Landscape)",
                                "9:16 (Portrait)",
                                "4:3 (Landscape)",
                                "3:4 (Portrait)",
                                "3:2 (Landscape)",
                                "2:3 (Portrait)",
                            ],
                            value="16:9 (Landscape)",
                            label="Output Aspect Ratio"
                        )

                        with gr.Row():
                            scene_steps = gr.Number(
                                label="Inference Steps",
                                value=4,
                                minimum=1,
                                maximum=50,
                                step=1
                            )
                            scene_guidance = gr.Number(
                                label="Guidance Scale",
                                value=1.0,
                                minimum=0.0,
                                maximum=10.0,
                                step=0.1
                            )

                        render_btn = gr.Button(
                            "RENDER SCENE",
                            variant="primary",
                            size="lg",
                            elem_classes=["generate-btn-main"]
                        )

                    # Right column: Output
                    with gr.Column(scale=2):
                        gr.Markdown("### Rendered Scene")

                        scene_output = gr.Image(
                            label="Scene Output",
                            type="pil",
                            format="png",
                            elem_classes=["output-image"]
                        )

                        scene_status = gr.Textbox(
                            label="Status",
                            interactive=False
                        )

                gr.Markdown("---")
                gr.Markdown("""
                **Tips for Scene Composer:**
                - Upload a character sheet generated in the first tab, or use any character turnaround image
                - Add a second character sheet to include multiple characters in the scene
                - Background images help set the scene location and lighting
                - Object/prop images can be items the character holds or interacts with
                - Be descriptive in your scene description for best results
                """)

        # Event handlers for Tab 1
        input_type.change(
            fn=update_input_visibility,
            inputs=[input_type],
            outputs=[main_input, face_input, body_input]
        )

        backend_dropdown.change(
            fn=update_defaults_on_backend_change,
            inputs=[backend_dropdown],
            outputs=[num_steps, guidance_scale, include_costume_in_faces]
        )

        generate_btn.click(
            fn=generate_character_sheet,
            inputs=[
                main_input,
                input_type,
                character_name,
                gender,
                costume_description,
                costume_image,
                face_input,
                body_input,
                backend_dropdown,
                api_key_input,
                num_steps,
                guidance_scale,
                include_costume_in_faces
            ],
            outputs=[
                output_image,
                status_text,
                preview_left_face,
                preview_front_face,
                preview_right_face,
                preview_left_body,
                preview_front_body,
                preview_right_body,
                preview_back_body,
                json_download,
                zip_download
            ]
        )

        # Event handlers for Tab 2 (Scene Composer)
        render_btn.click(
            fn=render_scene,
            inputs=[
                scene_char1,
                scene_char2,
                scene_background,
                scene_object,
                scene_description,
                scene_aspect_ratio,
                backend_dropdown,
                api_key_input,
                scene_steps,
                scene_guidance
            ],
            outputs=[
                scene_output,
                scene_status
            ]
        )

    return demo


# =============================================================================
# Main
# =============================================================================

if __name__ == "__main__":
    demo = create_ui()

    if HF_SPACES:
        # Running on HuggingFace Spaces
        demo.launch(
            theme=gr.themes.Soft(),
            css=APP_CSS
        )
    else:
        # Local testing
        print("Running locally (no Zero GPU)")
        demo.launch(
            server_name="0.0.0.0",
            server_port=7890,
            share=False,
            theme=gr.themes.Soft(),
            css=APP_CSS
        )