app.py

import gradio as gr
from gradio_pdf import PDF
import numpy as np
import random
import torch
import spaces
import math
import os
import yaml
import io
import tempfile
import shutil
import uuid
import time
import json
from typing import List, Tuple, Dict, Optional
from datetime import datetime, timedelta
from pathlib import Path

from PIL import Image
from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler
from huggingface_hub import InferenceClient
from reportlab.lib.pagesizes import A4
from reportlab.pdfgen import canvas
from reportlab.pdfbase import pdfmetrics
from reportlab.lib.utils import ImageReader
from PyPDF2 import PdfReader, PdfWriter

# --- Style Presets Loading ---

def load_style_presets():
    """Load style presets from YAML file."""
    try:
        with open('style_presets.yaml', 'r') as f:
            data = yaml.safe_load(f)
        # Filter only enabled presets
        presets = {k: v for k, v in data['presets'].items() if v.get('enabled', True)}
        return presets
    except Exception as e:
        print(f"Error loading style presets: {e}")
        return {"no_style": {"id": "no_style", "label": "No style (custom)", "prompt_prefix": "", "prompt_suffix": "", "negative_prompt": ""}}

# Load presets at startup
STYLE_PRESETS = load_style_presets()

# --- Page Layouts Loading ---

def load_page_layouts():
    """Load page layouts from YAML file."""
    try:
        with open('page_layouts.yaml', 'r') as f:
            data = yaml.safe_load(f)
        return data['layouts']
    except Exception as e:
        print(f"Error loading page layouts: {e}")
        # Fallback to basic layouts
        return {
            "1_image": [{"id": "full_page", "label": "Full Page", "positions": [[0.05, 0.05, 0.9, 0.9]]}],
            "2_images": [{"id": "horizontal_split", "label": "Horizontal Split", "positions": [[0.05, 0.05, 0.425, 0.9], [0.525, 0.05, 0.425, 0.9]]}],
            "3_images": [{"id": "grid", "label": "Grid", "positions": [[0.05, 0.05, 0.283, 0.5], [0.358, 0.05, 0.283, 0.5], [0.666, 0.05, 0.283, 0.5]]}],
            "4_images": [{"id": "grid_2x2", "label": "2x2 Grid", "positions": [[0.05, 0.05, 0.425, 0.425], [0.525, 0.05, 0.425, 0.425], [0.05, 0.525, 0.425, 0.425], [0.525, 0.525, 0.425, 0.425]]}]
        }

# Load layouts at startup
PAGE_LAYOUTS = load_page_layouts()

def get_layout_choices(num_images: int) -> List[Tuple[str, str]]:
    """Get available layout choices for a given number of images."""
    key = f"{num_images}_image" if num_images == 1 else f"{num_images}_images"
    if key in PAGE_LAYOUTS:
        return [(layout["label"], layout["id"]) for layout in PAGE_LAYOUTS[key]]
    # Return empty list if no layouts found (shouldn't happen with our config)
    return [("Default", "default")]

def get_layout_metadata(layout_id: str, num_images: int) -> List[Dict]:
    """Get metadata for each panel in a layout.
    
    Args:
        layout_id: ID of the selected layout
        num_images: Total number of images
        
    Returns:
        List of metadata dicts with panel_type, focus, composition, shot_type, and camera_angle
    """
    key = f"{num_images}_image" if num_images == 1 else f"{num_images}_images"
    layouts = PAGE_LAYOUTS.get(key, [])
    layout = next((l for l in layouts if l["id"] == layout_id), None)
    
    if layout and "metadata" in layout:
        return layout["metadata"]
    
    # Fallback metadata if not found
    fallback_meta = {
        "panel_type": "action",
        "focus": "character",
        "composition": "square",
        "shot_type": "medium",
        "camera_angle": "eye_level"
    }
    return [fallback_meta] * num_images

def get_random_style_preset():
    """Get a random style preset (excluding 'no_style' and 'random')."""
    eligible_keys = [k for k in STYLE_PRESETS.keys() if k not in ['no_style', 'random']]
    if eligible_keys:
        return random.choice(eligible_keys)
    return 'no_style'

def apply_style_preset(prompt, style_preset_key, custom_style_text=""):
    """
    Apply style preset to the prompt.

    Args:
        prompt: The user's base prompt
        style_preset_key: The key of the selected style preset
        custom_style_text: Custom style text when 'no_style' is selected

    Returns:
        tuple: (styled_prompt, negative_prompt)
    """
    if style_preset_key == 'no_style':
        # Use custom style text if provided
        if custom_style_text and custom_style_text.strip():
            styled_prompt = f"{custom_style_text}, {prompt}"
        else:
            styled_prompt = prompt
        return styled_prompt, ""

    if style_preset_key == 'random':
        # Select a random style
        style_preset_key = get_random_style_preset()

    if style_preset_key in STYLE_PRESETS:
        preset = STYLE_PRESETS[style_preset_key]
        prefix = preset.get('prompt_prefix', '')
        suffix = preset.get('prompt_suffix', '')
        negative = preset.get('negative_prompt', '')

        # Build the styled prompt
        parts = []
        if prefix:
            parts.append(prefix)
        parts.append(prompt)
        if suffix:
            parts.append(suffix)

        styled_prompt = ', '.join(parts)
        return styled_prompt, negative

    # Fallback to original prompt if preset not found
    return prompt, ""

# --- Story Generation using Hugging Face InferenceClient ---

def generate_story_scenes(story_prompt, num_scenes, style_context="", panel_metadata=None):
    """
    Generates a sequence of scene descriptions with captions and dialogues.

    Args:
        story_prompt: The user's story prompt
        num_scenes: Number of scenes to generate
        style_context: Optional style context to consider
        panel_metadata: List of metadata dicts for each panel

    Returns:
        List of dicts with 'caption' and 'dialogue' keys
    """
    # Ensure HF_TOKEN is set
    api_key = os.environ.get("HF_TOKEN")
    if not api_key:
        print("HF_TOKEN not set, using fallback scene generation")
        # Simple fallback - just split the prompt into scenes
        fallback_scenes = []
        for i in range(num_scenes):
            fallback_scenes.append({
                "caption": f"{story_prompt} (scene {i+1} of {num_scenes})",
                "dialogue": ""
            })
        return fallback_scenes

    # Initialize the client
    client = InferenceClient(
        provider="cerebras",
        api_key=api_key,
    )

    # Build panel descriptions from metadata
    panel_descriptions = []
    if panel_metadata and len(panel_metadata) == num_scenes:
        for i, meta in enumerate(panel_metadata):
            panel_type = meta.get('panel_type', 'action')
            focus = meta.get('focus', 'character')
            composition = meta.get('composition', 'square')
            shot_type = meta.get('shot_type', 'medium')
            camera_angle = meta.get('camera_angle', 'eye_level')
            
            # Format shot type for readability
            shot_display = shot_type.replace('_', ' ').title()
            angle_display = camera_angle.replace('_', ' ').title()
            
            # Create a descriptive text for this panel
            desc = f"Panel {i+1}/{num_scenes} - {composition.upper()} composition, {shot_display} shot at {angle_display} angle, {panel_type} panel focusing on {focus}"
            panel_descriptions.append(desc)
    else:
        # Fallback if no metadata
        panel_descriptions = [f"Panel {i+1}/{num_scenes}" for i in range(num_scenes)]

    # Create system prompt with panel-specific guidance
    system_prompt = f"""You are a story writer with expertise in cinematography and visual storytelling. Generate exactly {num_scenes} panels for a short story, based on the user's story prompt.

PANEL LAYOUT INFORMATION:
The page has {num_scenes} panels with the following characteristics:
{chr(10).join(f"- {desc}" for desc in panel_descriptions)}

IMPORTANT INSTRUCTIONS:
1. Output ONLY a YAML list with exactly {num_scenes} items
2. Each item must have exactly two fields:
   - caption: A detailed visual description of the scene (describe characters, clothing, location, action, expressions)
   - dialogue: Natural language description of what the character says/exclaims/shouts (can be empty string if no dialogue)

3. **ADAPT EACH SCENE TO ITS PANEL TYPE:**
   - ESTABLISHING panels: Describe the full environment, setting, atmosphere, time of day, location details
   - ACTION panels: Focus on dynamic movement, motion lines, impact, energy, physical activity
   - CLOSEUP panels: Describe facial features, eyes, expressions, emotions in extreme detail
   - DIALOGUE panels: Focus on character interactions, body language during conversation
   - REACTION panels: Emphasize emotional responses, facial expressions, body language
   - DETAIL panels: Zoom in on specific objects, hands, symbols, or small but important elements
   - TRANSITION panels: Show passage of time, change of location, or connecting moments
   - SPLASH panels: Epic, dramatic, full-scene moments with maximum visual impact

4. **ADAPT TO SHOT TYPE (CINEMATOGRAPHY):**
   - EXTREME WIDE SHOT: Vast landscapes, tiny characters in massive environments, epic scale
   - WIDE SHOT: Full scene with characters and environment, establishing context
   - FULL SHOT: Entire character from head to toe, showing their full body and stance
   - MEDIUM FULL SHOT: Character from knees up, showing most of body with some environment
   - MEDIUM SHOT: Character from waist up, balanced between character and setting
   - MEDIUM CLOSEUP: Head and shoulders, focusing on face while showing some context
   - CLOSEUP: Face filling frame, detailed facial features and expressions
   - EXTREME CLOSEUP: Tiny detail - just eyes, hands, mouth, or specific object filling frame

5. **ADAPT TO CAMERA ANGLE:**
   - EYE LEVEL: Neutral, straightforward angle - camera at character's eye level
   - HIGH ANGLE: Camera looking down on subject - can make them seem vulnerable, small, or overwhelmed
   - LOW ANGLE: Camera looking up at subject - makes them seem powerful, imposing, heroic
   - OVERHEAD/BIRD'S EYE: Camera directly above looking down - shows spatial relationships, isolation
   - DUTCH ANGLE/CANTED: Tilted camera - creates tension, disorientation, chaos, unease
   - OVER THE SHOULDER (OTS): Camera behind one character looking at another - intimate conversation
   - POV (Point of View): Camera as character's eyes - immersive, first-person perspective

6. **ADAPT TO COMPOSITION:**
   - WIDE/LANDSCAPE: Emphasize horizontal elements, panoramic views, sweeping scenes, breadth
   - TALL/PORTRAIT: Emphasize vertical elements, full-body shots, top-to-bottom action, height
   - SQUARE: Balanced composition, centered subjects, symmetrical arrangements

7. **ADAPT TO FOCUS:**
   - CHARACTER: Detailed character description (appearance, clothing, pose, expression)
   - CHARACTERS (plural): Multiple people, their relationships, positioning, interactions
   - ENVIRONMENT: Setting details, location, atmosphere, background elements, mood
   - EVENT: What's happening, the action, the moment being captured, the incident
   - EMOTION: Facial expression, body language, emotional state, feelings
   - OBJECT: Detailed description of an important item, prop, symbol, or artifact
   - ACTION: Movement, impact, dynamic poses, energy, motion

8. **CONSIDER PANEL PROGRESSION:**
   - You're creating panel X of {num_scenes} - consider where you are in the story flow
   - Early panels (1-2/{num_scenes}): Establish setting and introduce characters
   - Middle panels: Build action, develop conflict, show character reactions
   - Later panels ({num_scenes-1}-{num_scenes}/{num_scenes}): Resolve the moment, provide reaction or cliffhanger

9. For captions: Be VERY descriptive. Include shot type language like "wide shot of...", "close-up on...", "overhead view of...". Repeat character descriptions in each scene if needed.

10. For dialogue: Write as natural language action: "The [character] says: [what they say]" or "The [character] exclaims: [what they exclaim]"
   - DO NOT include character names in the dialogue text itself
   - Use verbs like: says, exclaims, shouts, whispers, asks, replies, thinks, mutters, screams

11. Keep continuity between scenes to tell a coherent story
12. Make each scene visually distinct but connected to the narrative

Example output format:
- caption: "Extreme wide shot from high angle of a dark alley at night, rain pouring down heavily, neon signs casting red and blue reflections in vast puddles covering the ground, tall buildings looming menacingly on both sides creating a narrow canyon, a young woman with long red hair wearing a blue detective coat stands small in the center of the frame examining glowing footprints on the wet pavement with a magnifying glass, dramatic lighting from above"
  dialogue: "The detective whispers to herself: These tracks... they're not human"
- caption: "Extreme close-up at eye level of the detective's piercing green eyes widening in shock and fear, her pupils dilating rapidly, individual beads of rain clinging to her dark eyelashes, her face illuminated by an eerie pulsing blue glow from below, wrinkles forming on her forehead"
  dialogue: ""
- caption: "Full shot at low angle, the red-haired detective in the blue coat leaping backwards dynamically with motion blur streaks, her coat billowing dramatically, a massive jagged shark fin erupting violently from a puddle behind her, water exploding upward in huge spray with droplets frozen mid-air, her expression one of pure terror, arms flailing"
  dialogue: "The detective shouts at the top of her lungs: OH NO! SHARKS IN THE CITY!"

Generate exactly {num_scenes} scenes. Output ONLY the YAML list, no other text."""

    # Format the messages
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": f"Create {num_scenes} scenes for this story: {story_prompt}"}
    ]

    try:
        # Call the API
        completion = client.chat.completions.create(
            model="Qwen/Qwen3-235B-A22B-Instruct-2507",
            messages=messages,
            temperature=0.7,
            max_tokens=2000,
        )
        response = completion.choices[0].message.content

        # Parse the YAML response
        scenes = parse_yaml_scenes(response, num_scenes)
        return scenes

    except Exception as e:
        print(f"Error during story generation: {e}")
        # Fallback to simple scene splitting
        fallback_scenes = []
        for i in range(num_scenes):
            fallback_scenes.append({
                "caption": f"{story_prompt} (part {i+1} of {num_scenes})",
                "dialogue": ""
            })
        return fallback_scenes

def parse_yaml_scenes(yaml_text, expected_count):
    """
    Parse YAML text to extract scene captions and dialogues.
    """
    try:
        # Clean up the text - remove markdown code blocks if present
        yaml_text = yaml_text.strip()
        if yaml_text.startswith("```yaml"):
            yaml_text = yaml_text[7:]
        if yaml_text.startswith("```"):
            yaml_text = yaml_text[3:]
        if yaml_text.endswith("```"):
            yaml_text = yaml_text[:-3]

        # Parse YAML
        scenes = yaml.safe_load(yaml_text)

        if not isinstance(scenes, list):
            raise ValueError("Expected a list of scenes")

        # Validate and clean scenes
        valid_scenes = []
        for scene in scenes:
            if isinstance(scene, dict) and 'caption' in scene:
                valid_scenes.append({
                    'caption': str(scene.get('caption', '')),
                    'dialogue': str(scene.get('dialogue', ''))
                })

        # Ensure we have the expected number of scenes
        while len(valid_scenes) < expected_count:
            valid_scenes.append({
                'caption': 'continuation of the story',
                'dialogue': ''
            })

        return valid_scenes[:expected_count]

    except Exception as e:
        print(f"Error parsing YAML scenes: {e}")
        # Return fallback scenes
        return [{'caption': 'scene description', 'dialogue': ''} for _ in range(expected_count)]


def get_caption_language(prompt):
    """Detects if the prompt contains Chinese characters."""
    ranges = [
        ('\u4e00', '\u9fff'),  # CJK Unified Ideographs
    ]
    for char in prompt:
        if any(start <= char <= end for start, end in ranges):
            return 'zh'
    return 'en'


# --- Model Loading ---
# Use the new lightning-fast model setup
ckpt_id = "Qwen/Qwen-Image"

# Scheduler configuration from the Qwen-Image-Lightning repository
scheduler_config = {
    "base_image_seq_len": 256,
    "base_shift": math.log(3),
    "invert_sigmas": False,
    "max_image_seq_len": 8192,
    "max_shift": math.log(3),
    "num_train_timesteps": 1000,
    "shift": 1.0,
    "shift_terminal": None,
    "stochastic_sampling": False,
    "time_shift_type": "exponential",
    "use_beta_sigmas": False,
    "use_dynamic_shifting": True,
    "use_exponential_sigmas": False,
    "use_karras_sigmas": False,
}

scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)
pipe = DiffusionPipeline.from_pretrained(
    ckpt_id, scheduler=scheduler, torch_dtype=torch.bfloat16
).to("cuda")

# Load LoRA weights for acceleration
pipe.load_lora_weights(
    "lightx2v/Qwen-Image-Lightning", weight_name="Qwen-Image-Lightning-8steps-V1.1.safetensors"
)
pipe.fuse_lora()

# --- UI Constants and Helpers ---
MAX_SEED = np.iinfo(np.int32).max

def get_image_size_for_position(position_data, image_index, num_images, max_resolution=1024):
    """
    Calculate EXACT image dimensions to match layout aspect ratio perfectly.
    
    This function calculates pixel dimensions that precisely match the aspect ratio
    of the layout rectangle to ensure images fill their containers without floating.
    
    Args:
        position_data: Layout position data [x, y, width, height] in relative units (0-1)
        image_index: Index of the current image (0-based)
        num_images: Total number of images in the layout
        max_resolution: Maximum resolution for any dimension (default 1024)

    Returns:
        tuple: (width, height) with exact aspect ratio matching the layout
    """
    if not position_data:
        return max_resolution, max_resolution  # Default square

    x_rel, y_rel, w_rel, h_rel = position_data

    # Calculate the EXACT aspect ratio from the layout rectangle
    # This is crucial - we must match this aspect ratio precisely
    # A4 page dimensions in points (must match PDF generation)
    page_width, page_height = 595.27, 841.89
    # Account for the page's aspect ratio when calculating layout aspect ratio
    layout_aspect_ratio = (w_rel / h_rel) * (page_width / page_height) if h_rel > 0 else 1.0

    # Scale to max_resolution while maintaining EXACT aspect ratio
    if layout_aspect_ratio >= 1:  # Wider than tall
        width = max_resolution
        height = max_resolution / layout_aspect_ratio
    else:  # Taller than wide
        height = max_resolution
        width = max_resolution * layout_aspect_ratio
    
    # Round to nearest 8 pixels for model compatibility
    # Using 8px instead of 64px preserves aspect ratio much better
    # Most diffusion models work well with multiples of 8
    width = round(width / 8) * 8
    height = round(height / 8) * 8
    
    # After rounding, ensure we maintain the aspect ratio as closely as possible
    # and don't exceed max_resolution
    if width > max_resolution:
        width = max_resolution
        height = round((max_resolution / layout_aspect_ratio) / 8) * 8
    if height > max_resolution:
        height = max_resolution
        width = round((max_resolution * layout_aspect_ratio) / 8) * 8
    
    # Ensure minimum size of 256px (reduced from 384 for more flexibility)
    # while maintaining the layout aspect ratio
    min_size = 256
    if width < min_size or height < min_size:
        if layout_aspect_ratio >= 1:  # Wider image
            width = max(min_size, width)
            height = round((width / layout_aspect_ratio) / 8) * 8
        else:  # Taller image
            height = max(min_size, height)
            width = round((height * layout_aspect_ratio) / 8) * 8
    
    # Final safety checks
    width = max(min_size, min(int(width), max_resolution))
    height = max(min_size, min(int(height), max_resolution))
    
    return width, height

def get_layout_position_for_image(layout_id, num_images, image_index):
    """Get the position data for a specific image in a layout.

    Args:
        layout_id: ID of the selected layout
        num_images: Total number of images
        image_index: Index of the current image (0-based)

    Returns:
        Position data [x, y, width, height] or None
    """
    key = f"{num_images}_image" if num_images == 1 else f"{num_images}_images"
    layouts = PAGE_LAYOUTS.get(key, [])
    layout = next((l for l in layouts if l["id"] == layout_id), None)

    if layout and "positions" in layout:
        positions = layout["positions"]
        if image_index < len(positions):
            return positions[image_index]

    # Fallback positions for each number of images
    fallback_positions = {
        1: [[0.05, 0.05, 0.9, 0.9]],
        2: [[0.05, 0.05, 0.425, 0.9], [0.525, 0.05, 0.425, 0.9]],
        3: [[0.05, 0.25, 0.283, 0.5], [0.358, 0.25, 0.283, 0.5], [0.666, 0.25, 0.283, 0.5]],
        4: [[0.05, 0.05, 0.425, 0.425], [0.525, 0.05, 0.425, 0.425],
            [0.05, 0.525, 0.425, 0.425], [0.525, 0.525, 0.425, 0.425]],
        5: [[0.05, 0.05, 0.9, 0.3], [0.05, 0.4, 0.283, 0.55], [0.358, 0.4, 0.283, 0.55],
            [0.666, 0.4, 0.283, 0.275], [0.666, 0.7, 0.283, 0.275]],
        6: [[0.05, 0.05, 0.425, 0.283], [0.525, 0.05, 0.425, 0.283],
            [0.05, 0.358, 0.425, 0.283], [0.525, 0.358, 0.425, 0.283],
            [0.05, 0.666, 0.425, 0.283], [0.525, 0.666, 0.425, 0.283]],
        7: [[0.28, 0.02, 0.44, 0.3], [0.02, 0.25, 0.3, 0.25], [0.68, 0.25, 0.3, 0.25],
            [0.25, 0.35, 0.5, 0.3], [0.02, 0.52, 0.3, 0.25], [0.68, 0.52, 0.3, 0.25],
            [0.28, 0.68, 0.44, 0.3]],
        8: [[0.02, 0.02, 0.23, 0.47], [0.27, 0.02, 0.23, 0.47], [0.52, 0.02, 0.23, 0.47],
            [0.77, 0.02, 0.21, 0.47], [0.02, 0.51, 0.23, 0.47], [0.27, 0.51, 0.23, 0.47],
            [0.52, 0.51, 0.23, 0.47], [0.77, 0.51, 0.21, 0.47]]
    }

    positions = fallback_positions.get(num_images, fallback_positions[1])
    if image_index < len(positions):
        return positions[image_index]
    return [0.05, 0.05, 0.9, 0.9]  # Ultimate default

# --- Session Management Functions ---

class SessionManager:
    """Manages user session data and temporary file storage."""

    def __init__(self, session_id: str = None):
        self.session_id = session_id or str(uuid.uuid4())
        self.base_dir = Path(tempfile.gettempdir()) / "gradio_comic_sessions"
        self.session_dir = self.base_dir / self.session_id
        self.session_dir.mkdir(parents=True, exist_ok=True)
        self.metadata_file = self.session_dir / "metadata.json"
        self.pdf_path = self.session_dir / "comic.pdf"
        self.load_or_create_metadata()

    def load_or_create_metadata(self):
        """Load existing metadata or create new."""
        if self.metadata_file.exists():
            with open(self.metadata_file, 'r') as f:
                self.metadata = json.load(f)
        else:
            self.metadata = {
                "created_at": datetime.now().isoformat(),
                "pages": [],
                "total_pages": 0
            }
            self.save_metadata()

    def save_metadata(self):
        """Save metadata to file."""
        with open(self.metadata_file, 'w') as f:
            json.dump(self.metadata, f, indent=2)

    def add_page(self, images: List[Image.Image], layout_id: str, seeds: List[int]):
        """Add a new page to the session."""
        page_num = self.metadata["total_pages"] + 1
        page_dir = self.session_dir / f"page_{page_num}"
        page_dir.mkdir(exist_ok=True)

        # Save images
        image_paths = []
        for i, img in enumerate(images):
            img_path = page_dir / f"image_{i+1}.jpg"
            img.save(img_path, 'JPEG', quality=95)
            image_paths.append(str(img_path))

        # Update metadata
        self.metadata["pages"].append({
            "page_num": page_num,
            "layout_id": layout_id,
            "num_images": len(images),
            "image_paths": image_paths,
            "seeds": seeds,
            "created_at": datetime.now().isoformat()
        })
        self.metadata["total_pages"] = page_num
        self.save_metadata()

        return page_num

    def get_all_pages_images(self) -> List[Tuple[List[Image.Image], str, int]]:
        """Get all images from all pages."""
        pages_data = []
        for page in self.metadata["pages"]:
            images = []
            for img_path in page["image_paths"]:
                if Path(img_path).exists():
                    images.append(Image.open(img_path))
            if images:
                pages_data.append((images, page["layout_id"], page["num_images"]))
        return pages_data

    def cleanup_old_sessions(self, max_age_hours: int = 24):
        """Clean up sessions older than max_age_hours."""
        if not self.base_dir.exists():
            return

        cutoff_time = datetime.now() - timedelta(hours=max_age_hours)

        for session_dir in self.base_dir.iterdir():
            if session_dir.is_dir():
                metadata_file = session_dir / "metadata.json"
                if metadata_file.exists():
                    try:
                        with open(metadata_file, 'r') as f:
                            metadata = json.load(f)
                        created_at = datetime.fromisoformat(metadata["created_at"])
                        if created_at < cutoff_time:
                            shutil.rmtree(session_dir)
                            print(f"Cleaned up old session: {session_dir.name}")
                    except Exception as e:
                        print(f"Error cleaning session {session_dir.name}: {e}")

# --- PDF Generation Functions ---

def create_single_page_pdf(images: List[Image.Image], layout_id: str, num_images: int) -> bytes:
    """
    Create a single PDF page with images arranged according to the selected layout.

    Args:
        images: List of PIL images
        layout_id: ID of the selected layout
        num_images: Number of images to include

    Returns:
        PDF page as bytes
    """
    # Create a bytes buffer for the PDF
    pdf_buffer = io.BytesIO()

    # Create canvas with A4 size
    pdf = canvas.Canvas(pdf_buffer, pagesize=A4)
    page_width, page_height = A4

    # Get the layout configuration
    key = f"{num_images}_image" if num_images == 1 else f"{num_images}_images"
    layouts = PAGE_LAYOUTS.get(key, [])
    layout = next((l for l in layouts if l["id"] == layout_id), None)

    if not layout:
        # Fallback to default grid layout with proper spacing
        if num_images == 1:
            positions = [[0.02, 0.02, 0.96, 0.96]]
        elif num_images == 2:
            positions = [[0.02, 0.02, 0.47, 0.96], [0.51, 0.02, 0.47, 0.96]]
        elif num_images == 3:
            positions = [[0.02, 0.2, 0.31, 0.6], [0.345, 0.2, 0.31, 0.6], [0.67, 0.2, 0.31, 0.6]]
        elif num_images == 4:
            positions = [[0.02, 0.02, 0.47, 0.47], [0.51, 0.02, 0.47, 0.47],
                        [0.02, 0.51, 0.47, 0.47], [0.51, 0.51, 0.47, 0.47]]
        elif num_images == 5:
            positions = [[0.02, 0.02, 0.96, 0.44], [0.02, 0.48, 0.31, 0.5], [0.345, 0.48, 0.31, 0.5],
                        [0.67, 0.48, 0.31, 0.24], [0.67, 0.74, 0.31, 0.24]]
        elif num_images == 6:
            positions = [[0.02, 0.02, 0.47, 0.31], [0.51, 0.02, 0.47, 0.31],
                        [0.02, 0.345, 0.47, 0.31], [0.51, 0.345, 0.47, 0.31],
                        [0.02, 0.67, 0.47, 0.31], [0.51, 0.67, 0.47, 0.31]]
        elif num_images == 7:
            positions = [[0.28, 0.02, 0.44, 0.3], [0.02, 0.25, 0.3, 0.25], [0.68, 0.25, 0.3, 0.25],
                        [0.25, 0.35, 0.5, 0.3], [0.02, 0.52, 0.3, 0.25], [0.68, 0.52, 0.3, 0.25],
                        [0.28, 0.68, 0.44, 0.3]]
        elif num_images == 8:
            positions = [[0.02, 0.02, 0.23, 0.47], [0.27, 0.02, 0.23, 0.47], [0.52, 0.02, 0.23, 0.47],
                        [0.77, 0.02, 0.21, 0.47], [0.02, 0.51, 0.23, 0.47], [0.27, 0.51, 0.23, 0.47],
                        [0.52, 0.51, 0.23, 0.47], [0.77, 0.51, 0.21, 0.47]]
        else:
            positions = [[0.02, 0.02, 0.96, 0.96]]
    else:
        positions = layout["positions"]

    # Draw each image according to the layout
    for i, (image, pos) in enumerate(zip(images[:num_images], positions)):
        if i >= len(images):
            break

        x_rel, y_rel, w_rel, h_rel = pos

        # Add small padding between panels (1% of page dimensions)
        padding = 0.01

        # Apply padding to prevent images from touching edges
        if x_rel < padding:
            x_rel = padding
        if y_rel < padding:
            y_rel = padding
        if x_rel + w_rel > 1 - padding:
            w_rel = 1 - padding - x_rel
        if y_rel + h_rel > 1 - padding:
            h_rel = 1 - padding - y_rel

        # Convert relative positions to absolute positions
        # Note: In ReportLab, y=0 is at the bottom
        x = x_rel * page_width
        y = (1 - y_rel - h_rel) * page_height  # Flip Y coordinate
        width = w_rel * page_width
        height = h_rel * page_height

        # Calculate aspect ratios for comparison
        img_aspect = image.width / image.height
        layout_aspect = width / height
        aspect_diff = abs(img_aspect - layout_aspect) / layout_aspect
        
        # If aspect ratios match closely (within 2%), fill the space completely
        # Otherwise, preserve aspect ratio to avoid distortion
        if aspect_diff < 0.02:  # Less than 2% difference
            # Aspect ratios match well - fill the space completely
            actual_width = width
            actual_height = height
            actual_x = x
            actual_y = y
        else:
            # Significant aspect ratio difference - preserve it to avoid distortion
            if img_aspect > layout_aspect:
                # Image is wider than the layout space
                new_height = width / img_aspect
                y_offset = (height - new_height) / 2
                actual_width = width
                actual_height = new_height
                actual_x = x
                actual_y = y + y_offset
            else:
                # Image is taller than the layout space
                new_width = height * img_aspect
                x_offset = (width - new_width) / 2
                actual_width = new_width
                actual_height = height
                actual_x = x + x_offset
                actual_y = y

        # Convert PIL image to format suitable for ReportLab
        img_buffer = io.BytesIO()
        image.save(img_buffer, format='JPEG', quality=95)
        img_buffer.seek(0)

        # Draw the image on the PDF
        # When aspect ratios match (aspect_diff < 0.02), we fill completely
        # Otherwise we preserve aspect ratio to prevent distortion
        pdf.drawImage(ImageReader(img_buffer), actual_x, actual_y,
                     width=actual_width, height=actual_height,
                     preserveAspectRatio=(aspect_diff >= 0.02), mask='auto')

    # Save the PDF
    pdf.save()

    # Get the PDF bytes
    pdf_buffer.seek(0)
    pdf_bytes = pdf_buffer.read()

    return pdf_bytes

def create_multi_page_pdf(session_manager: SessionManager) -> str:
    """
    Create a multi-page PDF from all pages in the session.

    Args:
        session_manager: SessionManager instance with page data

    Returns:
        Path to the created PDF file
    """
    pages_data = session_manager.get_all_pages_images()

    if not pages_data:
        return None

    # Create PDF writer
    pdf_writer = PdfWriter()

    # Create each page
    for images, layout_id, num_images in pages_data:
        page_pdf_bytes = create_single_page_pdf(images, layout_id, num_images)

        # Read the single page PDF
        page_pdf_reader = PdfReader(io.BytesIO(page_pdf_bytes))

        # Add the page to the writer
        for page in page_pdf_reader.pages:
            pdf_writer.add_page(page)

    # Write to file with explicit flushing
    pdf_path = session_manager.pdf_path
    with open(pdf_path, 'wb') as f:
        pdf_writer.write(f)
        f.flush()  # Flush Python's internal buffer
        os.fsync(f.fileno())  # Ensure OS writes to disk

    # Small delay to ensure file system catches up
    time.sleep(0.1)

    return str(pdf_path)

# --- Main Inference Function (with session support) ---
@spaces.GPU(duration=240)  # Increased duration for up to 8 images
def infer_page(
    prompt,
    guidance_scale=1.0,
    num_inference_steps=8,
    style_preset="no_style",
    custom_style_text="",
    num_images=1,
    layout="default",
    max_resolution=1024,
    session_state=None,
    progress=gr.Progress(track_tqdm=True),
):
    """
    Generates images for a new page and adds them to the PDF.

    Args:
        prompt (str): The text prompt to generate images from.
        guidance_scale (float): Corresponds to `true_cfg_scale`.
        num_inference_steps (int): The number of denoising steps.
        style_preset (str): The key of the style preset to apply.
        custom_style_text (str): Custom style text when 'no_style' is selected.
        num_images (int): Number of images to generate (1-8).
        layout (str): The layout ID for arranging images in the PDF.
        max_resolution: Maximum resolution for any dimension.
        session_state: Current session state dictionary.
        progress (gr.Progress): A Gradio Progress object to track generation.

    Returns:
        tuple: Updated session state, PDF path, and updated button label.
    """
    # Initialize or retrieve session
    if session_state is None or "session_id" not in session_state:
        session_state = {"session_id": str(uuid.uuid4()), "page_count": 0}

    session_manager = SessionManager(session_state["session_id"])

    # Clean up old sessions periodically
    if random.random() < 0.1:  # 10% chance to cleanup on each request
        session_manager.cleanup_old_sessions()

    # Check page limit (reduced to 24 for performance)
    if session_manager.metadata["total_pages"] >= 24:
        return session_state, None, f"Page limit reached"

    generated_images = []
    used_seeds = []

    # Get panel metadata for this layout
    panel_metadata = get_layout_metadata(layout, int(num_images))
    
    # Debug: print metadata
    print(f"\n=== LAYOUT METADATA for {layout} with {num_images} images ===")
    for i, meta in enumerate(panel_metadata):
        shot_type = meta.get('shot_type', 'medium').replace('_', ' ').title()
        camera_angle = meta.get('camera_angle', 'eye_level').replace('_', ' ').title()
        print(f"Panel {i+1}/{num_images}: {meta['panel_type']} | Focus: {meta['focus']} | {meta['composition']} | {shot_type} @ {camera_angle}")
    print("=" * 80 + "\n")

    # Generate story scenes with metadata
    progress(0, f"Generating story with {num_images} scenes...")
    scenes = generate_story_scenes(prompt, int(num_images), style_preset, panel_metadata)

    # Generate the requested number of images
    for i in range(int(num_images)):
        progress((i + 0.5) / num_images, f"Generating image {i+1} of {num_images} for page {session_manager.metadata['total_pages'] + 1}")

        current_seed = random.randint(0, MAX_SEED)  # Always randomize seed

        # Get optimal aspect ratio based on position in layout
        position_data = get_layout_position_for_image(layout, int(num_images), i)

        # Use scene caption and dialogue for this image
        scene_prompt = scenes[i]['caption']
        scene_dialogue = scenes[i]['dialogue']
        
        # Get metadata for this panel
        panel_meta = panel_metadata[i] if i < len(panel_metadata) else {}
        
        # Debug output
        print(f"\n--- Generating Panel {i+1}/{num_images} ---")
        print(f"Type: {panel_meta.get('panel_type', 'unknown')}")
        print(f"Focus: {panel_meta.get('focus', 'unknown')}")
        print(f"Composition: {panel_meta.get('composition', 'unknown')}")
        shot_display = panel_meta.get('shot_type', 'medium').replace('_', ' ').title()
        angle_display = panel_meta.get('camera_angle', 'eye_level').replace('_', ' ').title()
        print(f"Shot: {shot_display} at {angle_display} angle")
        print(f"Caption: {scene_prompt[:100]}..." if len(scene_prompt) > 100 else f"Caption: {scene_prompt}")
        print(f"Dialogue: {scene_dialogue if scene_dialogue else '(none)'}")
        print("-" * 80)

        # Generate single image with automatic aspect ratio
        image, used_seed = infer_single_auto(
            prompt=scene_prompt,
            seed=current_seed,
            randomize_seed=False,  # We handle randomization here
            position_data=position_data,
            image_index=i,
            num_images=int(num_images),
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            dialogue=scene_dialogue,  # Pass dialogue separately
            style_preset=style_preset,
            custom_style_text=custom_style_text,
            max_resolution=max_resolution,
        )

        generated_images.append(image)
        used_seeds.append(used_seed)

    # Add page to session
    progress(0.8, "Adding page to document...")
    page_num = session_manager.add_page(generated_images, layout, used_seeds)

    # Create multi-page PDF
    progress(0.9, "Creating PDF...")
    pdf_path = create_multi_page_pdf(session_manager)

    progress(1.0, "Done!")

    # Update session state
    session_state["page_count"] = page_num

    # Next button label
    next_page_num = page_num + 1
    button_label = f"Generate page {next_page_num}" if next_page_num <= 24 else "Page limit reached"

    return session_state, pdf_path, button_label

# New inference function with automatic aspect ratio
def infer_single_auto(
    prompt,
    seed=42,
    randomize_seed=False,
    position_data=None,
    image_index=0,
    num_images=1,
    guidance_scale=1.0,
    num_inference_steps=8,
    dialogue="",
    style_preset="no_style",
    custom_style_text="",
    max_resolution=1024,
):
    """
    Generates an image with automatically determined aspect ratio based on layout position.
    """
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    # Automatically determine image size based on position with custom max resolution
    width, height = get_image_size_for_position(position_data, image_index, num_images, max_resolution)
    
    # Calculate layout aspect ratio for verification
    if position_data:
        x_rel, y_rel, w_rel, h_rel = position_data
        layout_aspect = w_rel / h_rel if h_rel > 0 else 1.0
        image_aspect = width / height
        aspect_error = abs(image_aspect - layout_aspect) / layout_aspect * 100
        print(f"Image {image_index + 1}/{num_images}: Layout aspect={layout_aspect:.4f}, Image aspect={image_aspect:.4f}, Error={aspect_error:.2f}%")
    
    # Set up the generator for reproducibility
    generator = torch.Generator(device="cuda").manual_seed(seed)

    print(f"Original prompt: '{prompt}'")
    print(f"Style preset: '{style_preset}'")
    print(f"Auto-selected size based on layout: {width}x{height}")

    # Apply style preset first
    styled_prompt, style_negative_prompt = apply_style_preset(prompt, style_preset, custom_style_text)

    # Add dialogue to the prompt if present
    if dialogue and dialogue.strip():
        styled_prompt = f"{styled_prompt}. {dialogue.strip()}"

    # Use style negative prompt if available, otherwise default
    negative_prompt = style_negative_prompt if style_negative_prompt else " "

    print(f"Final Prompt: '{styled_prompt}'")
    print(f"Negative Prompt: '{negative_prompt}'")
    print(f"Seed: {seed}, Size: {width}x{height}, Steps: {num_inference_steps}, True CFG Scale: {guidance_scale}")

    # Generate the image
    image = pipe(
        prompt=styled_prompt,
        negative_prompt=negative_prompt,
        width=width,
        height=height,
        num_inference_steps=num_inference_steps,
        generator=generator,
        true_cfg_scale=guidance_scale,
    ).images[0]

    # Convert to grayscale if using manga_no_color style
    if style_preset == "manga_no_color":
        image = image.convert('L').convert('RGB')

    return image, seed

# Keep the old infer function for backward compatibility (simplified)
infer = infer_single_auto

# --- Examples and UI Layout ---
examples = [
        "A capybara wearing a suit holding a sign that reads Hello World",
]

css = """
#col-container {
    margin: 0 auto;
    max-width: 1024px;
}
#logo-title {
    text-align: center;
}
#logo-title img {
    width: 400px;
}
"""

with gr.Blocks(css=css) as demo:

    gr.set_static_paths(paths=["logo.png"])

    # Session state
    session_state = gr.State(value={"session_id": str(uuid.uuid4()), "page_count": 0})

    with gr.Column(elem_id="col-container"):

        # Row 1: Logo, Inputs, Buttons
        with gr.Row():
            # Column 1: Logo
            with gr.Column(scale=0, min_width=140):
                gr.HTML("""
                <div id="logo-title" style="display: flex; align-items: center; justify-content: center; padding: 0px 0; position: relative;">
                    <img src="/gradio_api/file=logo.png" alt="AI Comic Factory Logo" style="max-height: 100px; width: auto; margin-right: 0px; border-radius: 6px;">
                    <div style="position: relative; display: flex; flex-direction: column; align-items: center;">

                        <style>
                            @keyframes pulse {
                                0%, 100% {
                                    transform: scale(1) rotate(var(--rotation));
                                }
                                50% {
                                    transform: scale(1.05) rotate(var(--rotation));
                                }
                            }
                            .comic-label-top {
                                --rotation: -5deg;
                            }
                            .comic-label-bottom {
                                --rotation: 7deg;
                            }
                            .comic-label:hover {
                                transform: scale(1.1) rotate(var(--rotation)) !important;
                                transition: transform 0.3s ease;
                            }
                            .comic-label::before {
                                content: '';
                                position: absolute;
                                top: -8px;
                                right: -8px;
                                width: 30px;
                                height: 30px;
                                background: radial-gradient(circle, #ffeb3b, #ffc107);
                                border-radius: 50%;
                                border: 2px solid #fff;
                                box-shadow: 0 2px 8px rgba(255, 193, 7, 0.6);
                                z-index: -1;
                                animation: sparkle 1.5s ease-in-out infinite;
                            }
                            @keyframes sparkle {
                                0%, 100% {
                                    transform: scale(1);
                                    opacity: 1;
                                }
                                50% {
                                    transform: scale(1.2);
                                    opacity: 0.8;
                                }
                            }
                        </style>
                    </div>
                </div>
                """)

            # Column 2: Inputs
            with gr.Column(scale=1):
                # Sub-row 1: Prompt
                prompt = gr.Text(
                    label="Prompt",
                    show_label=False,
                    placeholder="Describe the current page",
                    container=False,
                )

                # Sub-row 2: Style, number of images, layout
                with gr.Row():
                    # Create dropdown choices from loaded presets
                    style_choices = [(preset["label"], key) for key, preset in STYLE_PRESETS.items()]
                    style_preset = gr.Dropdown(
                        label="Style Preset",
                        choices=style_choices,
                        value="no_style",
                        interactive=True,
                        scale=1
                    )

                    # Number of images slider
                    num_images_slider = gr.Slider(
                        label="Nb of panels",
                        minimum=1,
                        # we can support 8, but we will need better layouts
                        # also starting from 8 I notice some latency for story generation
                        maximum=7,
                        step=1,
                        value=6,
                        scale=1
                    )

                    # Page layout dropdown - initialize with correct choices for default value (6 images)
                    default_num_images = 6
                    default_layout_choices = get_layout_choices(default_num_images)
                    layout_dropdown = gr.Dropdown(
                        label="Page Layout",
                        choices=default_layout_choices,
                        value=default_layout_choices[0][1] if default_layout_choices else "default",
                        interactive=True,
                        scale=1
                    )

            # Column 3: Buttons
            with gr.Column(scale=0,min_width=200):
                run_button = gr.Button("Generate page 1", variant="primary")
                reset_button = gr.Button("Reset", variant="secondary")

        # Row 2: Settings column and PDF preview
        with gr.Row():
            # Left column - Additional settings
            with gr.Column(scale=0):
                custom_style_text = gr.Textbox(
                    label="Custom Style Text",
                    placeholder="Enter custom style (e.g., 'oil painting')",
                    visible=False,
                    lines=1
                )

                # Advanced settings accordion
                with gr.Accordion("Advanced Settings", open=False):
                    guidance_scale = gr.Slider(
                        label="Guidance scale (True CFG Scale)",
                        minimum=1.0,
                        maximum=5.0,
                        step=0.1,
                        value=1.0,
                    )

                    num_inference_steps = gr.Slider(
                        label="Number of inference steps",
                        minimum=4,
                        maximum=28,
                        step=1,
                        value=8,
                    )

                    max_resolution = gr.Slider(
                        label="Max Resolution",
                        minimum=768,
                        maximum=1280,
                        step=128,
                        value=1024,
                        info="Maximum dimension for generated images (higher = better quality but slower)"
                    )

                gr.Markdown("""**Note:** Your images and PDF are saved for up to 24 hours.
                You can continue adding pages (up to 24) by clicking the generate button.""")

            # Right column - PDF Preview
            with gr.Column(scale=2):
                pdf_preview = PDF(
                    label="PDF Preview",
                    show_label=True,
                    height=900,
                    elem_id="pdf-preview",
                    enable_zoom=True,
                    min_zoom=1.0,
                    max_zoom=3.0
                )

        # Add interaction to show/hide custom style text field
        def toggle_custom_style(style_value):
            return gr.update(visible=(style_value == "no_style"))

        style_preset.change(
            fn=toggle_custom_style,
            inputs=[style_preset],
            outputs=[custom_style_text]
        )

        # Update layout dropdown when number of images changes
        def update_layout_choices(num_images):
            choices = get_layout_choices(int(num_images))
            return gr.update(choices=choices, value=choices[0][1] if choices else "default")

        num_images_slider.change(
            fn=update_layout_choices,
            inputs=[num_images_slider],
            outputs=[layout_dropdown]
        )

    # Define the main generation event
    generation_event = gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer_page,
        inputs=[
            prompt,
            guidance_scale,
            num_inference_steps,
            style_preset,
            custom_style_text,
            num_images_slider,
            layout_dropdown,
            max_resolution,
            session_state,
        ],
        outputs=[session_state, pdf_preview, run_button],
    )

    # Reset button functionality
    def reset_session():
        new_state = {"session_id": str(uuid.uuid4()), "page_count": 0}
        return new_state, None, None, "Generate page 1"

    # Connect the reset button
    reset_button.click(
        fn=reset_session,
        inputs=[],
        outputs=[session_state, pdf_preview, run_button]
    )

if __name__ == "__main__":
    demo.launch(mcp_server=True)