import gradio as gr
from google import genai
from google.genai import types
from PIL import Image
import os
from pathlib import Path
from typing import Optional, Tuple, List
from io import BytesIO

# Beautiful CSS Design
custom_css = """
/* Modern Typography */
@import url('https://fonts.googleapis.com/css2?family=Poppins:wght@300;400;500;600;700&display=swap');

* {
    font-family: 'Poppins', sans-serif !important;
}

/* Animated Gradient Background */
.gradio-container {
    background: linear-gradient(-45deg, #ee7752, #e73c7e, #23a6d5, #23d5ab);
    background-size: 400% 400%;
    animation: gradientBG 15s ease infinite;
    min-height: 100vh;
}

@keyframes gradientBG {
    0% { background-position: 0% 50%; }
    50% { background-position: 100% 50%; }
    100% { background-position: 0% 50%; }
}

/* Main Container */
.main-wrap {
    max-width: 1400px !important;
    margin: 0 auto !important;
    padding: 20px !important;
}

/* Header Styling */
.header-container {
    background: rgba(255, 255, 255, 0.95);
    backdrop-filter: blur(10px);
    border-radius: 30px;
    padding: 3rem 2rem;
    text-align: center;
    margin-bottom: 2rem;
    box-shadow: 0 20px 60px rgba(0,0,0,0.15);
}

.header-title {
    font-size: 3.5rem;
    font-weight: 700;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    margin: 0 0 0.5rem 0;
}

.header-subtitle {
    font-size: 1.2rem;
    color: #64748b;
    font-weight: 400;
}

/* Card Styling */
.card {
    background: rgba(255, 255, 255, 0.95);
    backdrop-filter: blur(10px);
    border-radius: 20px;
    padding: 2rem;
    margin-bottom: 1.5rem;
    box-shadow: 0 10px 40px rgba(0,0,0,0.1);
    border: 1px solid rgba(255,255,255,0.5);
    width: 100%;
    box-sizing: border-box;
}

/* Gallery Styling */
.gallery-container {
    background: rgba(255, 255, 255, 0.95);
    border-radius: 20px;
    padding: 1.5rem;
    box-shadow: 0 10px 40px rgba(0,0,0,0.1);
}

/* Gallery sizing overrides */
#furniture_gallery {
    width: 100% !important;
}

#furniture_gallery .grid-wrap {
    display: grid !important;
    grid-template-columns: repeat(auto-fill, minmax(150px, 1fr)) !important;
    gap: 15px !important;
    width: 100% !important;
}

#furniture_gallery .grid-container {
    border-radius: 15px !important;
    overflow: hidden !important;
    transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1) !important;
    box-shadow: 0 4px 15px rgba(0,0,0,0.1) !important;
}

#furniture_gallery .grid-container:hover {
    transform: translateY(-5px) scale(1.03) !important;
    box-shadow: 0 10px 30px rgba(0,0,0,0.2) !important;
}

#furniture_gallery .grid-container.selected {
    border: 3px solid #23d5ab !important;
    box-shadow: 0 0 30px rgba(35,213,171,0.4) !important;
}

/* Button Styling */
.apply-btn {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    color: white !important;
    font-size: 1.1rem !important;
    font-weight: 600 !important;
    padding: 1rem 3rem !important;
    border-radius: 15px !important;
    border: none !important;
    cursor: pointer !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 5px 20px rgba(102,126,234,0.3) !important;
    text-transform: uppercase !important;
    letter-spacing: 1px !important;
    margin: 1.5rem auto !important;
    display: block !important;
}

.apply-btn:hover {
    transform: translateY(-3px) !important;
    box-shadow: 0 10px 30px rgba(102,126,234,0.4) !important;
}

/* Section Headers */
.section-header {
    font-size: 1.8rem;
    font-weight: 600;
    color: #1e293b;
    margin-bottom: 1.5rem;
    padding-left: 1rem;
    border-left: 5px solid #667eea;
}

/* Image Components */
.image-frame {
    border-radius: 15px !important;
    overflow: hidden !important;
    box-shadow: 0 10px 30px rgba(0,0,0,0.1) !important;
}

/* Status Messages */
.status-box {
    padding: 1rem 1.5rem;
    border-radius: 12px;
    font-weight: 500;
    margin: 1rem 0;
    text-align: center;
}

.status-success {
    background: linear-gradient(135deg, #84fab0 0%, #8fd3f4 100%);
    color: #065f46;
}

.status-error {
    background: linear-gradient(135deg, #fecaca 0%, #fbbf24 100%);
    color: #991b1b;
}

/* Tabs Styling */
.tabs {
    background: rgba(255, 255, 255, 0.95) !important;
    border-radius: 20px !important;
    padding: 1rem !important;
}

/* Feature Cards */
.feature-grid {
    display: grid;
    grid-template-columns: repeat(auto-fit, minmax(250px, 1fr));
    gap: 1.5rem;
    margin: 2rem 0;
}

.feature-card {
    background: linear-gradient(135deg, #f0f9ff 0%, #e0e7ff 100%);
    padding: 1.5rem;
    border-radius: 15px;
    text-align: center;
    transition: transform 0.3s ease;
}

.feature-card:hover {
    transform: translateY(-5px);
}

.feature-icon {
    font-size: 2.5rem;
    margin-bottom: 0.5rem;
}

.feature-title {
    font-weight: 600;
    color: #1e293b;
    margin-bottom: 0.5rem;
}

/* Responsive Design */
@media (max-width: 768px) {
    .header-title {
        font-size: 2.5rem;
    }
    
    #furniture_gallery .grid-wrap {
        grid-template-columns: repeat(auto-fill, minmax(100px, 1fr)) !important;
    }
    
    .apply-btn {
        padding: 0.8rem 2rem !important;
        font-size: 1rem !important;
    }
}

@media (max-width: 480px) {
    .header-title {
        font-size: 2rem;
    }
    
    .header-container {
        padding: 2rem 1rem;
    }
}

/* Loading Animation */
.loading {
    opacity: 0.6;
    animation: pulse 2s infinite;
}

@keyframes pulse {
    0%, 100% { opacity: 0.6; }
    50% { opacity: 1; }
}
"""

def initialize_client(api_key: str) -> Tuple[Optional[genai.Client], Optional[str]]:
    """Initialize the Gemini AI client"""
    try:
        client = genai.Client(api_key=api_key)
        return client, None
    except Exception as e:
        return None, f"❌ Error: {str(e)}"

def load_gallery_images() -> List[str]:
    """Load all furniture images from gallery folder"""
    gallery_path = Path("gallery")
    if not gallery_path.exists():
        os.makedirs(gallery_path, exist_ok=True)
        return []
    
    valid_extensions = {'.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp'}
    images = []
    
    for file in sorted(gallery_path.iterdir()):
        if file.suffix.lower() in valid_extensions:
            images.append(str(file.resolve()))
    
    return images

def get_image_from_gallery(index: int) -> Optional[Image.Image]:
    """Get a specific image from gallery by index"""
    images = load_gallery_images()
    if 0 <= index < len(images):
        try:
            return Image.open(images[index])
        except:
            return None
    return None

def analyze_and_place_furniture(
    client: genai.Client,
    room_image: Image.Image,
    furniture_image: Image.Image,
    placement_mode: str,
    custom_instructions: str = ""
) -> Tuple[Optional[Image.Image], str]:
    """Analyze room and place furniture using AI"""
    
    # First, analyze the room
    room_prompt = """
    Analyze this room image:
    1. Identify all structural elements (windows, doors, walls)
    2. Note existing furniture and their positions
    3. Identify empty spaces suitable for furniture
    4. Describe lighting and perspective
    """
    
    try:
        room_analysis = client.models.generate_content(
            model="gemini-2.5-flash-image",
            contents=[room_image, room_prompt]
        ).text
    except:
        room_analysis = "Room with furniture placement areas"
    
    # Identify furniture type
    furniture_prompt = """
    What type of furniture is this? 
    Describe its style, color, material, and any patterns.
    Is it a sofa, chair, table, curtain, rug, or something else?
    """
    
    try:
        furniture_desc = client.models.generate_content(
            model="gemini-2.5-flash-image",
            contents=[furniture_image, furniture_prompt]
        ).text
    except:
        furniture_desc = "Furniture item"
    
    # Check for special furniture types
    is_curtain = any(word in furniture_desc.lower() for word in ['curtain', 'drape', 'blind'])
    is_rug = any(word in furniture_desc.lower() for word in ['rug', 'carpet', 'mat'])
    
    # Build placement prompt
    placement_prompt = f"""
    You are an expert interior designer. Place this furniture in the room.
    
    ROOM ANALYSIS:
    {room_analysis}
    
    FURNITURE:
    {furniture_desc}
    
    PLACEMENT MODE: {placement_mode}
    {"- Replace similar existing furniture" if "Replace" in placement_mode else "- Add to empty space"}
    
    USER INSTRUCTIONS:
    {custom_instructions if custom_instructions else "Place naturally"}
    
    {"CURTAIN RULES: Place exactly on windows, maintain exact pattern and color" if is_curtain else ""}
    {"RUG RULES: Place flat on floor with correct perspective" if is_rug else ""}
    
    REQUIREMENTS:
    1. Keep room structure unchanged (walls, windows, doors)
    2. Maintain camera angle and lighting
    3. Ensure realistic scale and shadows
    4. Create photorealistic result
    5. Preserve furniture appearance exactly
    
    Generate a photorealistic image with the furniture integrated naturally.
    """
    
    try:
        # Generate the result
        response = client.models.generate_content(
            model="gemini-2.5-flash-image",
            contents=[room_image, furniture_image, placement_prompt],
            config=types.GenerateContentConfig(
                temperature=0.2,
                top_p=0.95
            )
        )
        
        # Extract image from response
        for part in response.candidates[0].content.parts:
            if part.inline_data:
                result_image = Image.open(BytesIO(part.inline_data.data))
                return result_image, "✅ Success! Furniture placed in your room."
        
        return None, "❌ No image generated. Please try again."
        
    except Exception as e:
        return None, f"❌ Error: {str(e)}"

def process_design_request(
    api_key: str,
    room_image: Image.Image,
    selected_furniture,
    placement_mode: str,
    instructions: str
) -> Tuple[Optional[Image.Image], str]:
    """Main processing function"""
    
    # Validate inputs
    if not room_image:
        return None, "⚠️ Please upload a room photo"
    
    if selected_furniture is None:
        return None, "⚠️ Please select furniture from the gallery"
    
    # Initialize client
    client, error = initialize_client(api_key)
    if error:
        return None, error
    
    # Get furniture image
    if isinstance(selected_furniture, list) and len(selected_furniture) > 0:
        furniture_path = selected_furniture[0]
    else:
        furniture_path = selected_furniture
    
    try:
        furniture_image = Image.open(furniture_path)
    except:
        return None, "❌ Error loading selected furniture"
    
    # Process the design
    return analyze_and_place_furniture(
        client, room_image, furniture_image, placement_mode, instructions
    )

def create_interface():
    """Create the main Gradio interface"""
    
    with gr.Blocks(
        theme=gr.themes.Soft(),
        css=custom_css,
        title="AI Interior Designer"
    ) as app:
        
        # Header
        gr.HTML("""
            <div class="header-container">
                <h1 class="header-title">AI Interior Designer</h1>
                <p class="header-subtitle">Transform your space with smart furniture placement</p>
            </div>
        """)
        
        with gr.Row():
            # Left Column - Controls
            with gr.Column(scale=1):
                with gr.Group(elem_classes=["card"]):
                    gr.Markdown("### 🔧 Configuration")
                    
                    api_key = gr.Textbox(
                        label="API Key",
                        type="password",
                        value="",
                        placeholder="Your Google AI API key"
                    )
                    
                    room_image = gr.Image(
                        label="Upload Room Photo",
                        type="pil",
                        elem_classes=["image-frame"]
                    )
                    
                    placement_mode = gr.Radio(
                        label="Placement Mode",
                        choices=["Add to Empty Space", "Replace Similar Furniture"],
                        value="Add to Empty Space"
                    )
                    
                    instructions = gr.Textbox(
                        label="Instructions (Optional)",
                        placeholder="E.g., Place sofa by window, Install curtains on main window",
                        lines=2
                    )
            
            # Right Column - Gallery
            with gr.Column(scale=2):
                with gr.Group(elem_classes=["card"]):
                    gr.Markdown("### 🛋️ Furniture Gallery")
                    gr.Markdown("*Select an item by clicking on it*")
                    
                    gallery_images = load_gallery_images()
                    
                    if gallery_images:
                        furniture_gallery = gr.Gallery(
                            value=gallery_images,
                            label="Available Furniture",
                            show_label=False,
                            elem_id="furniture_gallery",
                            columns=4,
                            rows=3,
                            height=400,
                            object_fit="cover",
                            interactive=True,
                            type="filepath",
                            preview=False
                        )
                    else:
                        gr.Markdown("⚠️ No images found in gallery folder. Please add images to the 'gallery' directory.")
                        furniture_gallery = gr.Gallery(
                            value=[],
                            visible=False,
                            type="filepath",
                            preview=False
                        )
                    
                    # Selected preview
                    selected_preview = gr.Image(
                        label="Selected Furniture",
                        type="pil",
                        height=180,
                        visible=False,
                        elem_classes=["image-frame"]
                    )
        
        # Apply Button
        apply_btn = gr.Button(
            "✨ Apply Furniture to Room",
            variant="primary",
            elem_classes=["apply-btn"]
        )
        
        # Output Section
        with gr.Group(elem_classes=["card"]):
            gr.Markdown("### 🎨 Result")
            
            output_image = gr.Image(
                label="Your Redesigned Room",
                type="pil",
                elem_classes=["image-frame"]
            )
            
            status = gr.HTML(
                value="",
                elem_classes=["status-box"]
            )
        
        # Features Section
        with gr.Accordion("✨ Features", open=False):
            gr.HTML("""
                <div class="feature-grid">
                    <div class="feature-card">
                        <div class="feature-icon">🎯</div>
                        <div class="feature-title">Smart Placement</div>
                        <div>AI understands room layout and places furniture naturally</div>
                    </div>
                    <div class="feature-card">
                        <div class="feature-icon">🪟</div>
                        <div class="feature-title">Window Detection</div>
                        <div>Automatically places curtains on detected windows</div>
                    </div>
                    <div class="feature-card">
                        <div class="feature-icon">💡</div>
                        <div class="feature-title">Realistic Lighting</div>
                        <div>Maintains proper shadows and reflections</div>
                    </div>
                    <div class="feature-card">
                        <div class="feature-icon">📐</div>
                        <div class="feature-title">Perfect Scale</div>
                        <div>Ensures furniture size matches room proportions</div>
                    </div>
                </div>
            """)
        
        # Tips
        with gr.Accordion("💡 Tips", open=False):
            gr.Markdown("""
                **For Best Results:**
                - Use clear, well-lit room photos
                - Take photos from corner angles
                - Select furniture that matches your room style
                - Be specific in instructions
                - Try different placement modes
                
                **Special Items:**
                - **Curtains**: Automatically placed on windows
                - **Rugs**: Positioned flat with correct perspective
                - **Sofas**: Placed against walls or in open spaces
                
                **Multiple Items:**
                - Run the tool multiple times
                - Build your room design step by step
            """)
        
        # Event Handlers
        selected_furniture_path = gr.State(value=None)
        gallery_images_state = gr.State(value=gallery_images)
        
        def gallery_select(evt: gr.SelectData, images):
            """Handle gallery selection"""
            selected_path = None

            if evt is not None:
                if evt.value:
                    if isinstance(evt.value, str):
                        selected_path = evt.value
                    elif isinstance(evt.value, list) and evt.value and isinstance(evt.value[0], str):
                        selected_path = evt.value[0]
                if selected_path is None and evt.index is not None and 0 <= evt.index < len(images):
                    selected_path = images[evt.index]

            if selected_path:
                try:
                    preview_img = Image.open(selected_path)
                    return selected_path, gr.update(value=preview_img, visible=True)
                except Exception:
                    return selected_path, gr.update(visible=False)

            return None, gr.update(visible=False)
        
        furniture_gallery.select(
            gallery_select,
            inputs=[gallery_images_state],
            outputs=[selected_furniture_path, selected_preview]
        )
        
        def apply_design(api_key, room_img, furniture_path, mode, inst):
            """Apply the selected furniture"""
            result_img, message = process_design_request(
                api_key, room_img, furniture_path, mode, inst
            )
            
            if result_img:
                status_html = f'<div class="status-success">{message}</div>'
            else:
                status_html = f'<div class="status-error">{message}</div>'
            
            return result_img, status_html
        
        apply_btn.click(
            apply_design,
            inputs=[
                api_key,
                room_image,
                selected_furniture_path,
                placement_mode,
                instructions
            ],
            outputs=[output_image, status]
        )
    
    return app

if __name__ == "__main__":
    app = create_interface()
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True
    )