app.py

import gradio as gr
from PIL import Image
import numpy as np
import time
import os
import random
from typing import List
import traceback
from google.oauth2.credentials import Credentials
from google_auth_oauthlib.flow import InstalledAppFlow
from googleapiclient.discovery import build
from googleapiclient.http import MediaIoBaseUpload
from io import BytesIO
import datetime

# Import the model interface
from model import DesignModel

# For testing, import the mock model
from mock_model import MockDesignModel

def create_ui(model: DesignModel):
    """Create the main UI interface with all components"""
    # Store current variations at UI level
    current_variations = []
    
    with gr.Blocks(css="""
        /* Base styles */
        :root {
            --section-title-size: 1.2rem;
            --section-spacing: var(--spacing-lg);
            --panel-min-height: auto;
        }

        /* Row styling for equal heights */
        .gr-row {
            margin-bottom: var(--spacing-md);
            display: flex;
            align-items: stretch;
        }

        .gr-row > .gr-column {
            display: flex;
            flex-direction: column;
        }

        .gr-row > .gr-column > .gr-group {
            flex: 1;
            display: flex;
            flex-direction: column;
        }

        /* Consistent title styling */
        .gr-markdown h2 {
            font-size: var(--section-title-size) !important;
            font-weight: 600 !important;
            margin: var(--spacing-sm) 0 var(--spacing-md) !important;
            padding: 0 !important;
            color: var(--body-text-color) !important;
        }

        .gr-markdown h3 {
            font-size: 1rem !important;
            font-weight: 500 !important;
            margin: var(--spacing-sm) 0 !important;
            padding: 0 !important;
            color: var(--body-text-color) !important;
        }

        /* Panel styling */
        .gr-group {
            border: 1px solid var(--border-color-primary);
            border-radius: var(--radius-lg);
            padding: var(--section-spacing);
            margin: var(--spacing-sm) 0;
            background: var(--background-fill-primary);
            min-height: var(--panel-min-height);
            height: auto !important;
            display: flex;
            flex-direction: column;
        }

        /* Form and input spacing */
        .gr-form {
            gap: var(--spacing-sm);
            flex-grow: 1;
            display: flex;
            flex-direction: column;
        }

        .gr-form > div {
            gap: var(--spacing-sm);
        }

        /* Dropdown styling */
        .gr-dropdown {
            margin-bottom: var(--spacing-sm);
        }

        /* Gallery improvements */
        #gallery {
            margin-top: 0;
            height: 300px !important;
        }

        #gallery img {
            object-fit: contain !important;
            width: 100% !important;
            height: 100% !important;
            max-height: none !important;
        }

        /* Button styling */
        .button-row {
            display: flex;
            justify-content: center;
            padding: var(--spacing-xl) 0;
        }

        .button-row button {
            min-width: 200px;
            font-size: 1.1em;
            font-weight: 600;
        }

        /* Text areas */
        .gr-textarea {
            font-family: monospace;
            line-height: 1.4;
        }

        .gr-textarea:disabled {
            opacity: 0.9;
            background-color: var(--background-fill-secondary);
        }

        /* Progress indicator */
        .progress-bar {
            margin: var(--spacing-sm) 0;
        }

        /* Make dropdowns always visible */
        .gr-dropdown {
            display: block !important;
            visibility: visible !important;
        }

        /* Upload area */
        .upload-group {
            height: 100%;
        }

        .upload-group .gr-image {
            min-height: 300px;
        }

        /* Checkbox alignment */
        .gr-checkbox-row {
            display: flex !important;
            align-items: center !important;
            min-height: 4.5rem !important;
        }

        .gr-checkbox-row .gr-checkbox {
            margin: auto 0 !important;
        }

        /* Remove any fixed heights from groups */
        .gr-group > div {
            height: auto !important;
            min-height: unset !important;
        }

        /* Ensure consistent spacing in all panels */
        .gr-group > div:not(:last-child) {
            margin-bottom: var(--spacing-sm);
        }

        /* Make surface finishes more compact */
        .surface-finishes {
            padding: var(--spacing-xs) !important;
            margin: 0 !important;
            border: 1px solid var(--border-color-primary) !important;
            border-radius: var(--radius-lg) !important;
            background: var(--background-fill-primary) !important;
        }
        
        .surface-finishes .gr-form {
            gap: var(--spacing-xs) !important;
            margin: 0 !important;
            padding: 0 !important;
            flex-grow: 0 !important;
        }

        .surface-finishes .gr-dropdown {
            margin: 0 0 var(--spacing-xs) 0 !important;
        }

        .surface-finishes .gr-row {
            margin: 0 !important;
            gap: var(--spacing-sm) !important;
        }

        .surface-finishes .gr-group {
            padding: 0 !important;
            margin: 0 !important;
            border: none !important;
            background: none !important;
            box-shadow: none !important;
            min-height: 0 !important;
        }

        .surface-finishes .gr-markdown {
            margin: 0 0 var(--spacing-xs) 0 !important;
        }

        .surface-finishes .gr-form > div {
            gap: var(--spacing-xs) !important;
            margin: 0 !important;
        }

        .surface-finishes .gr-column {
            flex-grow: 0 !important;
            padding: 0 !important;
        }

        /* Remove any minimum heights */
        .surface-finishes .gr-group > div {
            min-height: 0 !important;
            height: auto !important;
            margin: 0 !important;
        }

        /* Override any flex growth */
        .surface-finishes .gr-form,
        .surface-finishes .gr-column,
        .surface-finishes .gr-group {
            flex: 0 0 auto !important;
        }

        /* Wall decorations and special requests */
        .wall-decorations-row > .gr-column > .gr-group {
            height: 100%;
        }

        /* Upload and gallery row */
        .upload-gallery-row > .gr-column > .gr-group {
            height: 100%;
        }
    """) as interface:
        gr.Markdown("### Interior Design Assistant")
        
        with gr.Blocks():
            # Row 1 - Basic Settings
            with gr.Row():
                with gr.Group():
                    gr.Markdown("## 🏠 Basic Settings")
                    with gr.Row():
                        room_type = gr.Dropdown(
                            choices=[
                                "Living Room", "Bedroom", "Kitchen", "Dining Room",
                                "Bathroom", "Home Office", "Kids Room", "Master Bedroom",
                                "Guest Room", "Studio Apartment", "Entryway", "Hallway",
                                "Game Room", "Library", "Home Theater", "Gym"
                            ],
                            label="Room Type",
                            value="Living Room"
                        )
                        style_preset = gr.Dropdown(
                            choices=[
                                "Modern", "Contemporary", "Minimalist", "Industrial",
                                "Scandinavian", "Mid-Century Modern", "Traditional",
                                "Transitional", "Farmhouse", "Rustic", "Bohemian",
                                "Art Deco", "Coastal", "Mediterranean", "Japanese",
                                "French Country", "Victorian", "Colonial", "Gothic",
                                "Baroque", "Rococo", "Neoclassical", "Eclectic",
                                "Zen", "Tropical", "Shabby Chic", "Hollywood Regency",
                                "Southwestern", "Asian Fusion", "Retro"
                            ],
                            label="Design Style",
                            value="Modern"
                        )
                        color_scheme = gr.Dropdown(
                            choices=[
                                "Neutral", "Monochromatic", "Minimalist White",
                                "Warm Gray", "Cool Gray", "Earth Tones",
                                "Pastel", "Bold Primary", "Jewel Tones",
                                "Black and White", "Navy and Gold", "Forest Green",
                                "Desert Sand", "Ocean Blue", "Sunset Orange",
                                "Deep Purple", "Emerald Green", "Ruby Red",
                                "Sapphire Blue", "Golden Yellow", "Sage Green",
                                "Dusty Rose", "Charcoal", "Cream", "Burgundy",
                                "Teal", "Copper", "Silver", "Bronze", "Slate"
                            ],
                            label="Color Mood",
                            value="Neutral"
                        )

            # Row 2 - Surface Finishes
            with gr.Row():
                # Floor Options
                with gr.Column(scale=1):
                    with gr.Group():
                        gr.Markdown("## 🎨 Floor Options")
                        floor_type = gr.Dropdown(
                            choices=[
                                "Keep Existing", "Hardwood", "Stone Tiles", "Porcelain Tiles",
                                "Soft Carpet", "Polished Concrete", "Marble", "Vinyl",
                                "Natural Bamboo", "Cork", "Ceramic Tiles", "Terrazzo",
                                "Slate", "Travertine", "Laminate", "Engineered Wood",
                                "Mosaic Tiles", "Luxury Vinyl Tiles", "Stained Concrete"
                            ],
                            label="Material",
                            value="Keep Existing"
                        )
                        floor_color = gr.Dropdown(
                            choices=[
                                "Keep Existing", "Light Oak", "Rich Walnut", "Cool Gray",
                                "Whitewashed", "Warm Cherry", "Deep Brown", "Classic Black",
                                "Natural", "Sandy Beige", "Chocolate", "Espresso",
                                "Honey Oak", "Weathered Gray", "White Marble",
                                "Cream Travertine", "Dark Slate", "Golden Teak",
                                "Rustic Pine", "Ebony"
                            ],
                            label="Color",
                            value="Keep Existing"
                        )
                        floor_pattern = gr.Dropdown(
                            choices=[
                                "Keep Existing", "Classic Straight", "Elegant Herringbone",
                                "V-Pattern", "Decorative Parquet", "Diagonal Layout",
                                "Basketweave", "Chevron", "Random Length", "Grid Pattern",
                                "Versailles Pattern", "Running Bond", "Hexagonal",
                                "Moroccan Pattern", "Brick Layout", "Diamond Pattern",
                                "Windmill Pattern", "Large Format", "Mixed Width"
                            ],
                            label="Pattern",
                            value="Keep Existing"
                        )
                
                # Wall Options
                with gr.Column(scale=1):
                    with gr.Group():
                        gr.Markdown("## 🎨 Wall Options")
                        wall_type = gr.Dropdown(
                            choices=[
                                "Keep Existing", "Fresh Paint", "Designer Wallpaper",
                                "Textured Finish", "Wood Panels", "Exposed Brick",
                                "Natural Stone", "Wooden Planks", "Modern Concrete",
                                "Venetian Plaster", "Wainscoting", "Shiplap",
                                "3D Wall Panels", "Fabric Panels", "Metal Panels",
                                "Cork Wall", "Tile Feature", "Glass Panels",
                                "Acoustic Panels", "Living Wall"
                            ],
                            label="Treatment",
                            value="Keep Existing"
                        )
                        wall_color = gr.Dropdown(
                            choices=[
                                "Keep Existing", "Crisp White", "Soft White", "Warm Beige",
                                "Gentle Gray", "Sky Blue", "Nature Green", "Sunny Yellow",
                                "Blush Pink", "Deep Blue", "Bold Black", "Sage Green",
                                "Terracotta", "Navy Blue", "Charcoal Gray", "Lavender",
                                "Olive Green", "Dusty Rose", "Teal", "Burgundy"
                            ],
                            label="Color",
                            value="Keep Existing"
                        )
                        wall_finish = gr.Dropdown(
                            choices=[
                                "Keep Existing", "Soft Matte", "Subtle Eggshell",
                                "Pearl Satin", "Sleek Semi-Gloss", "High Gloss",
                                "Suede Texture", "Metallic", "Chalk Finish",
                                "Distressed", "Brushed", "Smooth", "Textured",
                                "Venetian", "Lime Wash", "Concrete", "Rustic",
                                "Lacquered", "Hammered", "Patina"
                            ],
                            label="Finish",
                            value="Keep Existing"
                        )

            # Row 3 - Wall Decorations and Special Requests
            with gr.Row(elem_classes="wall-decorations-row"):
                # Wall Decorations
                with gr.Column(scale=2):
                    with gr.Group():
                        gr.Markdown("## 🖼️ Wall Decorations")
                        # Art and Mirror
                        with gr.Row():
                            # Art Print
                            with gr.Column():
                                with gr.Row():
                                    art_print_enable = gr.Checkbox(label="Add Artwork", value=False)
                                art_print_color = gr.Dropdown(
                                    choices=[
                                        "None", "Classic Black & White", "Vibrant Colors",
                                        "Single Color", "Soft Colors", "Modern Abstract",
                                        "Earth Tones", "Pastel Palette", "Bold Primary Colors",
                                        "Metallic Accents", "Monochromatic", "Jewel Tones",
                                        "Watercolor", "Vintage Colors", "Neon Accents",
                                        "Natural Hues", "Ocean Colors", "Desert Palette"
                                    ],
                                    label="Art Style",
                                    value="None"
                                )
                                art_print_size = gr.Dropdown(
                                    choices=[
                                        "None", "Modest", "Standard", "Statement", "Oversized",
                                        "Gallery Wall", "Diptych", "Triptych", "Mini Series",
                                        "Floor to Ceiling", "Custom Size"
                                    ],
                                    label="Art Size",
                                    value="None"
                                )
                            
                            # Mirror
                            with gr.Column():
                                with gr.Row():
                                    mirror_enable = gr.Checkbox(label="Add Mirror", value=False)
                                mirror_frame = gr.Dropdown(
                                    choices=[
                                        "None", "Gold", "Silver", "Black", "White", "Wood",
                                        "Brass", "Bronze", "Copper", "Chrome", "Antique Gold",
                                        "Brushed Nickel", "Rustic Wood", "Ornate", "Minimalist",
                                        "LED Backlit", "Bamboo", "Rattan", "Leather Wrapped"
                                    ],
                                    label="Frame Style",
                                    value="None"
                                )
                                mirror_size = gr.Dropdown(
                                    choices=[
                                        "Small", "Medium", "Large", "Full Length",
                                        "Oversized", "Double Width", "Floor Mirror",
                                        "Vanity Size", "Statement Piece", "Custom Size"
                                    ],
                                    label="Mirror Size",
                                    value="Medium"
                                )

                        # Sconce, Shelf, and Plants
                        with gr.Row():
                            # Sconce
                            with gr.Column():
                                with gr.Row():
                                    sconce_enable = gr.Checkbox(label="Add Wall Sconce", value=False)
                                sconce_color = gr.Dropdown(
                                    choices=[
                                        "None", "Black", "Gold", "Silver", "Bronze", "White",
                                        "Brass", "Copper", "Chrome", "Antique Brass",
                                        "Brushed Nickel", "Oil-Rubbed Bronze", "Pewter",
                                        "Rose Gold", "Matte Black", "Polished Nickel",
                                        "Aged Brass", "Champagne", "Gunmetal"
                                    ],
                                    label="Sconce Color",
                                    value="None"
                                )
                                sconce_style = gr.Dropdown(
                                    choices=[
                                        "Modern", "Traditional", "Industrial", "Art Deco",
                                        "Minimalist", "Vintage", "Contemporary", "Rustic",
                                        "Coastal", "Farmhouse", "Mid-Century", "Bohemian",
                                        "Scandinavian", "Asian", "Mediterranean", "Gothic",
                                        "Transitional", "Eclectic", "Victorian"
                                    ],
                                    label="Sconce Style",
                                    value="Modern"
                                )

                            # Floating Shelves
                            with gr.Column():
                                with gr.Row():
                                    shelf_enable = gr.Checkbox(label="Add Floating Shelves", value=False)
                                shelf_color = gr.Dropdown(
                                    choices=[
                                        "None", "White", "Black", "Natural Wood", "Glass",
                                        "Dark Wood", "Light Wood", "Metal", "Gold", "Silver",
                                        "Bronze", "Reclaimed Wood", "Bamboo", "Marble",
                                        "Industrial Metal", "Two-Tone", "Concrete",
                                        "Acrylic", "Copper", "Brass"
                                    ],
                                    label="Shelf Material",
                                    value="None"
                                )
                                shelf_size = gr.Dropdown(
                                    choices=[
                                        "Small", "Medium", "Large", "Set of 3",
                                        "Extra Long", "Corner Set", "Asymmetric Set",
                                        "Graduated Sizes", "Custom Length", "Mini Cubes",
                                        "Full Wall", "Mixed Sizes", "Modular System"
                                    ],
                                    label="Shelf Size",
                                    value="Medium"
                                )

                            # Plants
                            with gr.Column():
                                with gr.Row():
                                    plants_enable = gr.Checkbox(label="Add Plants", value=False)
                                plants_type = gr.Dropdown(
                                    choices=[
                                        "None", "Hanging Plants", "Vertical Garden",
                                        "Plant Shelf", "Single Plant", "Climbing Vines",
                                        "Air Plants", "Succulent Wall", "Herb Garden",
                                        "Mixed Tropical", "Fern Collection", "Living Wall",
                                        "Moss Wall", "Potted Arrangement", "Plant Corner",
                                        "Cascading Plants", "Bamboo Screen", "Terrarium Wall"
                                    ],
                                    label="Plant Type",
                                    value="None"
                                )
                                plants_size = gr.Dropdown(
                                    choices=[
                                        "Small", "Medium", "Large", "Mixed Sizes",
                                        "Full Wall", "Statement Piece", "Compact",
                                        "Expansive", "Accent", "Floor to Ceiling",
                                        "Window Height", "Custom Size", "Modular"
                                    ],
                                    label="Plant Coverage",
                                    value="Medium"
                                )

                # Special Requests and Advanced Settings
                with gr.Column(scale=1):
                    with gr.Group():
                        gr.Markdown("## ✨ Special Requests")
                        input_text = gr.Textbox(
                            label="Additional Details",
                            placeholder="Add any special requests or details here...",
                            lines=3
                        )
                        num_outputs = gr.Slider(
                            minimum=1, maximum=50, value=1, step=1,
                            label="Number of Variations"
                        )
                        
                        gr.Markdown("### Advanced Settings")
                        num_steps = gr.Slider(
                            minimum=20,
                            maximum=100,
                            value=50,
                            step=1,
                            label="Quality Steps"
                        )
                        guidance_scale = gr.Slider(
                            minimum=1,
                            maximum=20,
                            value=7.5,
                            step=0.1,
                            label="Design Freedom"
                        )
                        strength = gr.Slider(
                            minimum=0.1,
                            maximum=1.0,
                            value=0.75,
                            step=0.05,
                            label="Change Amount"
                        )
                        seed = gr.Number(
                            label="Seed (leave empty for random)",
                            value=-1,
                            precision=0
                        )
                        with gr.Row():
                            save_to_drive = gr.Checkbox(label="Save to Google Drive")
                            drive_url = gr.Textbox(
                                label="Drive Folder URL",
                                placeholder="https://drive.google.com/drive/folders/..."
                            )

            # Row 4 - Current Prompts
            with gr.Row():
                with gr.Group():
                    gr.Markdown("## 📝 Current Prompts")
                    prompt_display = gr.TextArea(
                        label="Positive Prompt",
                        interactive=False,
                        lines=3,
                        value="Your design prompt will appear here..."
                    )
                    negative_prompt = gr.TextArea(
                        label="Negative Prompt",
                        value="blurry, low quality, distorted, deformed, disfigured, watermark, text, bad proportions, duplicate, double, multiple, broken, cropped",
                        lines=2,
                        interactive=False
                    )

            # Row 5 - Upload and Gallery
            with gr.Row(elem_classes="upload-gallery-row"):
                # Upload Area
                with gr.Column(scale=1):
                    with gr.Group():
                        gr.Markdown("## 📸 Upload Photo")
                        input_image = gr.Image(
                            label="Upload a photo of your room",
                            type='pil'
                        )
                
                # Gallery Area
                with gr.Column(scale=2):
                    with gr.Group():
                        gr.Markdown("## 🖼️ Generated Variations")
                        gallery = gr.Gallery(
                            show_label=False,
                            elem_id="gallery",
                            columns=4,
                            rows=1,
                            height="300px",
                            object_fit="contain",
                            preview=True,
                            show_share_button=False
                        )

            # Row 6 - Create Button
            with gr.Row(elem_classes="button-row"):
                submit = gr.Button("✨ Create My Design", variant="primary", size="lg")

            # Progress indicator
            progress = gr.Progress(track_tqdm=True)

        def upload_to_drive(image, folder_id):
            """Upload an image to Google Drive folder"""
            try:
                # OAuth 2.0 scopes
                SCOPES = ['https://www.googleapis.com/auth/drive.file']
                
                # Start OAuth 2.0 flow
                flow = InstalledAppFlow.from_client_secrets_file(
                    'credentials.json',
                    SCOPES
                )
                creds = flow.run_local_server(port=0)
                
                # Build the Drive API service
                service = build('drive', 'v3', credentials=creds)
                
                # Convert numpy array to bytes
                img = Image.fromarray(image)
                img_byte_arr = BytesIO()
                img.save(img_byte_arr, format='PNG')
                img_byte_arr.seek(0)
                
                # Prepare the file metadata
                timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
                file_metadata = {
                    'name': f'design_variation_{timestamp}.png',
                    'parents': [folder_id]
                }
                
                # Create media
                media = MediaIoBaseUpload(
                    img_byte_arr,
                    mimetype='image/png',
                    resumable=True
                )
                
                # Execute the upload
                file = service.files().create(
                    body=file_metadata,
                    media_body=media,
                    fields='id'
                ).execute()
                
                print(f"File uploaded successfully. File ID: {file.get('id')}")
                return True
                
            except Exception as e:
                print(f"Error uploading to Drive: {e}")
                return False

        def extract_folder_id(drive_url):
            """Extract folder ID from Google Drive URL"""
            try:
                # Handle different URL formats
                if 'folders/' in drive_url:
                    folder_id = drive_url.split('folders/')[-1].split('?')[0]
                    return folder_id
            except Exception as e:
                print(f"Error extracting folder ID: {e}")
            return None

        def on_submit(image, room, style, colors, floor_t, floor_c, floor_p,
                     wall_t, wall_c, wall_f, custom_text,
                     art_en, art_col, art_size,
                     mirror_en, mirror_fr, mirror_size,
                     sconce_en, sconce_col, sconce_style,
                     shelf_en, shelf_col, shelf_size,
                     plants_en, plants_type, plants_size,
                     num_outputs, save_to_drive, drive_url, num_steps, 
                     guidance_scale, seed, strength):
            
            if image is None:
                return []
            
            try:
                nonlocal current_variations
                
                # Generate the prompt
                prompt = update_prompt(
                    room, style, colors, floor_t, floor_c, floor_p,
                    wall_t, wall_c, wall_f, custom_text,
                    art_en, art_col, art_size,
                    mirror_en, mirror_fr, mirror_size,
                    sconce_en, sconce_col, sconce_style,
                    shelf_en, shelf_col, shelf_size,
                    plants_en, plants_type, plants_size
                )
                
                # Generate variations
                variations = model.generate_design(
                    image=image,
                    num_variations=max(1, int(num_outputs)),
                    prompt=prompt,
                    num_steps=int(num_steps),
                    guidance_scale=float(guidance_scale),
                    strength=float(strength),
                    seed=int(seed) if seed != -1 else None
                )
                
                # Store variations
                current_variations = variations
                
                # Handle Google Drive upload if enabled
                if save_to_drive and drive_url:
                    folder_id = extract_folder_id(drive_url)
                    if folder_id:
                        for variation in variations:
                            upload_to_drive(variation, folder_id)
                
                # Convert variations to gallery format
                gallery_images = [(v, None) for v in variations]
                return gallery_images
                
            except Exception as e:
                print(f"Error in generation: {e}")
                current_variations = []
                return []

        submit.click(
            on_submit,
            inputs=[
                input_image, room_type, style_preset, color_scheme,
                floor_type, floor_color, floor_pattern,
                wall_type, wall_color, wall_finish,
                input_text,
                art_print_enable, art_print_color, art_print_size,
                mirror_enable, mirror_frame, mirror_size,
                sconce_enable, sconce_color, sconce_style,
                shelf_enable, shelf_color, shelf_size,
                plants_enable, plants_type, plants_size,
                num_outputs, save_to_drive, drive_url, num_steps,
                guidance_scale, seed, strength
            ],
            outputs=[gallery]
        )

        # Update prompt display when any input changes
        def update_prompt_display(*args):
            try:
                prompt = update_prompt(*args)
                return [prompt, negative_prompt.value]  # Return both prompts
            except Exception as e:
                print(f"Error updating prompt: {e}")
                return ["Error generating prompt", negative_prompt.value]

        # List of all inputs that should trigger prompt updates
        prompt_inputs = [
            room_type, style_preset, color_scheme,
            floor_type, floor_color, floor_pattern,
            wall_type, wall_color, wall_finish,
            input_text,
            art_print_enable, art_print_color, art_print_size,
            mirror_enable, mirror_frame, mirror_size,
            sconce_enable, sconce_color, sconce_style,
            shelf_enable, shelf_color, shelf_size,
            plants_enable, plants_type, plants_size
        ]

        # Connect all inputs to prompt update
        for input_component in prompt_inputs:
            input_component.change(
                fn=update_prompt_display,
                inputs=prompt_inputs,
                outputs=[prompt_display, negative_prompt]
            )

        # Gallery click handler
        def on_select(evt):
            nonlocal current_variations
            try:
                if isinstance(evt, list) and len(evt) > 0:
                    # Get the clicked file path
                    clicked_path = evt[0][0] if isinstance(evt[0], tuple) else evt[0]
                    
                    # Get all file paths from the gallery
                    gallery_paths = []
                    for item in evt:
                        path = item[0] if isinstance(item, tuple) else item
                        gallery_paths.append(path)
                    
                    # Find which image was clicked by comparing paths
                    selected_index = gallery_paths.index(clicked_path)
                    if 0 <= selected_index < len(current_variations):
                        return current_variations[selected_index]
            except Exception as e:
                print(f"Gallery selection error: {e}")
            return None

        gallery.select(
            fn=on_select,
            inputs=gallery,
            outputs=[]
        )
    
    return interface

def update_prompt(room, style, colors, floor_t, floor_c, floor_p,
                 wall_t, wall_c, wall_f, custom_text,
                 art_en, art_col, art_size,
                 mirror_en, mirror_fr, mirror_size,
                 sconce_en, sconce_col, sconce_style,
                 shelf_en, shelf_col, shelf_size,
                 plants_en, plants_type, plants_size):
    # Start with basic room and style
    prompt_parts = [f"Design a {style} {room.lower()} with a {colors} color scheme"]
    
    # Add floor details only if type is not "Keep Existing"
    if floor_t and floor_t != "Keep Existing":
        floor_desc = floor_t
        if floor_c and floor_c != "Keep Existing":
            floor_desc += f" in {floor_c}"
        if floor_p and floor_p != "Keep Existing":
            floor_desc += f" with {floor_p} pattern"
        prompt_parts.append(f"featuring {floor_desc} flooring")
    
    # Add wall details only if type is not "Keep Existing"
    if wall_t and wall_t != "Keep Existing":
        wall_desc = wall_t
        if wall_c and wall_c != "Keep Existing":
            wall_desc += f" in {wall_c}"
        if wall_f and wall_f != "Keep Existing":
            wall_desc += f" with {wall_f} finish"
        prompt_parts.append(f"with {wall_desc} walls")
    
    # Add accessories only if enabled AND properties are selected and not "Keep Existing" or "None"
    accessories = []
    
    # Art Print
    if art_en and art_col and art_col not in ["Keep Existing", "None"]:
        accessories.append(f"{art_size} {art_col} Art Print")
    
    # Mirror
    if mirror_en and mirror_fr and mirror_fr not in ["Keep Existing", "None"]:
        accessories.append(f"{mirror_size} Mirror with {mirror_fr} frame")
    
    # Wall Sconce
    if sconce_en and sconce_col and sconce_col not in ["Keep Existing", "None"]:
        accessories.append(f"{sconce_style} {sconce_col} Wall Sconce")
    
    # Floating Shelves
    if shelf_en and shelf_col and shelf_col not in ["Keep Existing", "None"]:
        accessories.append(f"{shelf_size} {shelf_col} Floating Shelves")
    
    # Wall Plants
    if plants_en and plants_type and plants_type not in ["Keep Existing", "None"]:
        accessories.append(f"{plants_size} {plants_type}")
    
    # Only add accessories section if there are any accessories
    if accessories:
        prompt_parts.append("decorated with " + ", ".join(accessories))
    
    # Add custom text only if provided and non-empty
    if custom_text and custom_text.strip():
        prompt_parts.append(custom_text.strip())
    
    return ", ".join(prompt_parts)

def main():
    """Main entry point for the application"""
    import sys
    
    # Check if we're in test mode
    is_test_mode = "--test" in sys.argv
    
    if is_test_mode:
        print("Starting in TEST mode...")
        from mock_model import MockDesignModel
        model = MockDesignModel()
    else:
        print("Starting in PRODUCTION mode...")
        from prod_model import ProductionDesignModel
        model = ProductionDesignModel()
    
    interface = create_ui(model)
    interface.launch(
        share=False,
        show_api=False,  # Hide API docs
        show_error=True  # Show errors for debugging
    )

if __name__ == "__main__":
    main()