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 user interface for the application"""
    with gr.Blocks() as interface:
        gr.Markdown("## 🏠 Basic Settings")
        with gr.Row():
            with gr.Column():
                room_type = gr.Dropdown(
                    choices=[
                        "None", "Living Room", "Bedroom", "Dining Room", "Kitchen",
                        "Bathroom", "Home Office", "Master Bedroom", "Guest Room",
                        "Family Room", "Study", "Game Room", "Library",
                        "Nursery", "Craft Room"
                    ],
                    label="Room Type",
                    value="None"
                )
                style_preset = gr.Dropdown(
                    choices=[
                        "None", "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="None"
                )
                color_scheme = gr.Dropdown(
                    choices=[
                        "None", "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="None"
                )

        # 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=10.0,
                        step=0.1,
                        label="Design Freedom"
                    )
                    strength = gr.Slider(
                        minimum=0.1,
                        maximum=1.0,
                        value=0.9,
                        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")

        def is_valid_value(*values):
            """Helper function to check if values are valid (not None or Keep Existing)"""
            return all(v not in ["None", "Keep Existing"] for v in values)

        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):
            """Generate a prompt for the design, skipping any None or Keep Existing values."""
            
            invalid_values = {"None", "Keep Existing", None, ""}
            prompt_parts = []
            
            # Basic room settings - each part separate for maximum flexibility
            if all(x not in invalid_values for x in [style, room, colors]):
                base_parts = []
                if style not in invalid_values:
                    base_parts.append(style)
                if room not in invalid_values:
                    base_parts.append(room.lower())
                if colors not in invalid_values:
                    base_parts.append(f"with a {colors} color scheme")
                if base_parts:
                    prompt_parts.append("Design a " + " ".join(base_parts))
            
            # Floor details - build up piece by piece
            if floor_t not in invalid_values:
                floor_parts = [floor_t]
                if floor_c not in invalid_values:
                    floor_parts.append(f"in {floor_c}")
                if floor_p not in invalid_values:
                    floor_parts.append(f"with {floor_p} pattern")
                prompt_parts.append(f"featuring {' '.join(floor_parts)} flooring")
            
            # Wall details - build up piece by piece
            if wall_t not in invalid_values:
                wall_parts = [wall_t]
                if wall_c not in invalid_values:
                    wall_parts.append(f"in {wall_c}")
                if wall_f not in invalid_values:
                    wall_parts.append(f"with {wall_f} finish")
                prompt_parts.append(f"with {' '.join(wall_parts)} walls")
            
            # Accessories - only add valid combinations
            accessories = []
            
            # Art Print
            if art_en and art_col not in invalid_values and art_size not in invalid_values:
                accessories.append(f"{art_size} {art_col} Art Print")
            
            # Mirror
            if mirror_en and mirror_fr not in invalid_values and mirror_size not in invalid_values:
                accessories.append(f"{mirror_size} Mirror with {mirror_fr} frame")
            
            # Wall Sconce
            if sconce_en and sconce_col not in invalid_values and sconce_style not in invalid_values:
                accessories.append(f"{sconce_style} {sconce_col} Wall Sconce")
            
            # Floating Shelves
            if shelf_en and shelf_col not in invalid_values and shelf_size not in invalid_values:
                accessories.append(f"{shelf_size} {shelf_col} Floating Shelves")
            
            # Wall Plants
            if plants_en and plants_type not in invalid_values and plants_size not in invalid_values:
                accessories.append(f"{plants_size} {plants_type}")
            
            if accessories:
                prompt_parts.append("decorated with " + ", ".join(accessories))
            
            # Add custom text if present
            if custom_text and custom_text.strip() and custom_text not in invalid_values:
                prompt_parts.append(custom_text.strip())
            
            return ", ".join(prompt_parts) if prompt_parts else "Please select room type, style, and color scheme"

        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:
                # 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,
                    prompt=prompt,
                    num_variations=max(1, int(num_outputs)),
                    num_steps=int(num_steps),
                    guidance_scale=float(guidance_scale),
                    strength=float(strength),
                    seed=int(seed) if seed != -1 else None
                )
                
                # 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}")
                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]
            )

    return interface

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()