"""
UI Components for Email Rule Agent - Refactored with proper session management
Self-contained implementation with all necessary functions
"""

import gradio as gr
import os
import json
import uuid
from datetime import datetime
from typing import List, Dict, Any, Tuple

# Import UI utility functions
from .ui_utils import (
    create_email_html,
    create_interactive_rule_cards,
    create_preview_banner,
    filter_emails,
    get_folder_dropdown_choices,
    get_preview_rules_count,
    get_folder_counts,
    extract_folder_name,
    format_date,
    sort_emails
)

# Import agent
from .simple_agent import process_chat_message as process_agent_message
from .ui_tools import process_chat_with_tools
from .ui_chat import process_chat_streaming

# Import MCP client  
from .mcp_client import MCPClient




def handle_preview_rule(
    rule_id: str,
    sort_option: str,
    search_query: str,
    mcp_client: Any,
    pending_rules,  # Removed type hint to avoid Gradio API generation issue
    current_emails,  # Removed type hint to avoid Gradio API generation issue
    sample_emails,  # Removed type hint to avoid Gradio API generation issue
    preview_emails  # Removed type hint to avoid Gradio API generation issue
) -> Tuple:
    """
    Handle rule preview by delegating to MCP backend
    
    Returns all 7 outputs expected by the UI:
    1. folder_dropdown update
    2. email_display HTML
    3. rule_cards HTML
    4. preview_banner HTML
    5. status_msg
    6. pending_rules state
    7. current_emails state
    """
    
    print(f"DEBUG: handle_preview_rule called with rule_id='{rule_id}'")
    
    # Validate input
    if not rule_id or rule_id.strip() == "":
        print("DEBUG: No rule ID provided")
        return (
            gr.update(),  # No dropdown change
            create_email_html(current_emails),
            create_interactive_rule_cards(pending_rules),
            create_preview_banner(0),
            "❌ Error: No rule ID provided",
            pending_rules,
            current_emails
        )
    
    # Find the rule - use rule_id consistently
    print(f"DEBUG: Looking for rule with id='{rule_id}' in {len(pending_rules)} rules")
    print(f"DEBUG: Available rule IDs: {[r.get('rule_id', 'NO_ID') for r in pending_rules]}")
    
    rule = next((r for r in pending_rules if r.get("rule_id") == rule_id), None)
        
    if not rule:
        print(f"DEBUG: Rule not found with rule_id='{rule_id}'")
        return (
            gr.update(),
            create_email_html(current_emails),
            create_interactive_rule_cards(pending_rules),
            create_preview_banner(0),
            f"❌ Rule not found: {rule_id}",
            pending_rules,
            current_emails
        )
    
    try:
        # Call MCP backend to get preview
        preview_response = mcp_client.preview_rule(rule)
        
        if preview_response.get('success', False):
            # Create a new list of rules with updated status
            updated_rules = []
            for r in pending_rules:
                if r.get("rule_id") == rule_id:
                    # Create a new rule object with preview status
                    updated_rule = r.copy()
                    updated_rule["status"] = "preview"
                    updated_rules.append(updated_rule)
                else:
                    updated_rules.append(r)
            
            # Get the affected emails from backend
            affected_emails = preview_response.get('affected_emails', [])
            
            # If we got emails, use them for display
            if affected_emails:
                display_emails = affected_emails
            else:
                # Fallback to current emails
                display_emails = current_emails
            
            # Get folder choices from the emails we're displaying
            folder_choices = get_folder_dropdown_choices(display_emails)
            if not folder_choices:
                folder_choices = ["Inbox (0)"]
            
            # Filter and create display
            folder_value = folder_choices[0]
            filtered = filter_emails(folder_value, search_query, sort_option, display_emails)
            email_html = create_email_html(filtered, folder_value)
            
            # Create status message
            stats = preview_response.get('statistics', {})
            status_msg = f"👁️ Preview: {stats.get('matched_count', 0)} emails would be affected"
            
            return (
                gr.update(choices=folder_choices, value=folder_value),
                email_html,
                create_interactive_rule_cards(updated_rules),
                create_preview_banner(get_preview_rules_count(updated_rules)),
                status_msg,
                updated_rules,  # Return the new list of rules
                display_emails  # Update current emails to preview data
            )
        else:
            # Handle structured errors from the backend or client
            error_message = preview_response.get('error', 'Preview failed')
            status_code = preview_response.get('status_code')
            
            # Show specific error message
            status_msg = f"❌ {error_message}"
            return (
                gr.update(),  # No dropdown change
                create_email_html(current_emails),
                create_interactive_rule_cards(pending_rules),
                create_preview_banner(get_preview_rules_count(pending_rules)),
                status_msg,
                pending_rules,
                current_emails
            )
            
    except Exception as e:
        print(f"Error in preview handler: {e}")
        status_msg = f"❌ Connection error: {str(e)}"
        return (
            gr.update(),  # No dropdown change
            create_email_html(current_emails),
            create_interactive_rule_cards(pending_rules),
            create_preview_banner(get_preview_rules_count(pending_rules)),
            status_msg,
            pending_rules,
            current_emails
        )





def handle_accept_rule(rule_id, mcp_client, current_folder, sort_option, search_query,
                      pending_rules, applied_rules, current_emails, sample_emails):
    """Handle rule acceptance via MCP backend"""
    
    if not rule_id or rule_id.strip() == "":
        filtered_emails = filter_emails(current_folder, search_query, sort_option, current_emails)
        return (create_interactive_rule_cards(pending_rules), 
                create_preview_banner(get_preview_rules_count(pending_rules)), 
                "❌ Error: No rule ID provided",
                create_email_html(filtered_emails, current_folder),
                gr.update(choices=get_folder_dropdown_choices(current_emails)),
                pending_rules, applied_rules, current_emails, sample_emails)
    
    # Find the rule
    rule = next((r for r in pending_rules if r.get("rule_id") == rule_id), None)
    if not rule:
        filtered_emails = filter_emails(current_folder, search_query, sort_option, current_emails)
        return (create_interactive_rule_cards(pending_rules), 
                create_preview_banner(get_preview_rules_count(pending_rules)), 
                f"❌ Rule not found: {rule_id}",
                create_email_html(filtered_emails, current_folder),
                gr.update(),
                pending_rules, applied_rules, current_emails, sample_emails)
    
    try:
        # Apply rule via MCP
        apply_response = mcp_client.apply_rule(rule, preview=False)
        
        if apply_response.get('success', False):
            # Rule is already saved by the apply_rule endpoint, no need to save again
            
            # Re-fetch current state from backend (single source of truth)
            updated_rules = mcp_client.get_rules()
            
            # Get updated emails if provided
            updated_emails = apply_response.get('updated_emails', current_emails)
            
            # Update UI
            stats = apply_response.get('statistics', {})
            status_msg = f"✅ Rule applied: {stats.get('processed_count', 0)} emails processed"
            
            new_folder_choices = get_folder_dropdown_choices(updated_emails)
            filtered_emails = filter_emails(current_folder, search_query, sort_option, updated_emails)
            
            return (create_interactive_rule_cards(updated_rules), 
                    create_preview_banner(get_preview_rules_count(updated_rules)), 
                    status_msg,
                    create_email_html(filtered_emails, current_folder),
                    gr.update(choices=new_folder_choices),
                    updated_rules, applied_rules, updated_emails, updated_emails)
        else:
            error_msg = apply_response.get('error', 'Failed to apply rule')
            status_msg = f"❌ {error_msg}"
            filtered_emails = filter_emails(current_folder, search_query, sort_option, current_emails)
            return (create_interactive_rule_cards(pending_rules), 
                    create_preview_banner(get_preview_rules_count(pending_rules)), 
                    status_msg,
                    create_email_html(filtered_emails, current_folder),
                    gr.update(),
                    pending_rules, applied_rules, current_emails, sample_emails)
            
    except Exception as e:
        print(f"Error calling MCP apply: {e}")
        status_msg = f"❌ Connection error: {str(e)}"
        filtered_emails = filter_emails(current_folder, search_query, sort_option, current_emails)
        return (create_interactive_rule_cards(pending_rules), 
                create_preview_banner(get_preview_rules_count(pending_rules)), 
                status_msg,
                create_email_html(filtered_emails, current_folder),
                gr.update(),
                pending_rules, applied_rules, current_emails, sample_emails)






def handle_run_rule(rule_id, mcp_client, current_folder, sort_option, search_query,
                   pending_rules, applied_rules, current_emails, sample_emails):
    """Handle running an accepted rule via MCP backend"""
    
    if not rule_id or rule_id.strip() == "":
        filtered_emails = filter_emails(current_folder, search_query, sort_option, current_emails)
        return (create_interactive_rule_cards(pending_rules), 
                create_preview_banner(get_preview_rules_count(pending_rules)), 
                "❌ Error: No rule ID provided",
                create_email_html(filtered_emails, current_folder),
                gr.update(choices=get_folder_dropdown_choices(current_emails)),
                pending_rules, applied_rules, current_emails, sample_emails)
    
    for rule in pending_rules:
        if rule["id"] == rule_id and rule.get("status") == "accepted":
            try:
                # Apply rule via MCP client
                apply_response = mcp_client.apply_rule(rule, preview=False)
                
                if apply_response.get('success', False):
                    stats = apply_response.get('statistics', {})
                    status_msg = f"▶️ Rule executed: {stats.get('processed_count', 0)} emails processed"
                    
                    # Update emails if provided
                    if 'updated_emails' in apply_response:
                        sample_emails = apply_response['updated_emails']
                        current_emails = sample_emails
                    
                    # Update email display with new data
                    new_folder_choices = get_folder_dropdown_choices(current_emails)
                    filtered_emails = filter_emails(current_folder, search_query, sort_option, current_emails)
                    
                    print(f"Rule executed: {rule['name']} (ID: {rule_id})")
                    return (create_interactive_rule_cards(pending_rules), 
                            create_preview_banner(get_preview_rules_count(pending_rules)), 
                            status_msg,
                            create_email_html(filtered_emails, current_folder),
                            gr.update(choices=new_folder_choices),
                            pending_rules, applied_rules, current_emails, sample_emails)
                else:
                    error_message = apply_response.get('error', 'Failed to run rule')
                    status_msg = f"❌ {error_message}"
                    filtered_emails = filter_emails(current_folder, search_query, sort_option, current_emails)
                    return (create_interactive_rule_cards(pending_rules), 
                            create_preview_banner(get_preview_rules_count(pending_rules)), 
                            status_msg,
                            create_email_html(filtered_emails, current_folder),
                            gr.update(),
                            pending_rules, applied_rules, current_emails, sample_emails)
                    
            except Exception as e:
                print(f"Error running rule: {e}")
                status_msg = "⚠️ Connection issue - could not run rule"
                filtered_emails = filter_emails(current_folder, search_query, sort_option, current_emails)
                return (create_interactive_rule_cards(pending_rules), 
                        create_preview_banner(get_preview_rules_count(pending_rules)), 
                        status_msg,
                        create_email_html(filtered_emails, current_folder),
                        gr.update(),
                        pending_rules, applied_rules, current_emails, sample_emails)
    
    filtered_emails = filter_emails(current_folder, search_query, sort_option, current_emails)
    return (create_interactive_rule_cards(pending_rules), 
            create_preview_banner(get_preview_rules_count(pending_rules)), 
            f"❌ Rule not found or not accepted: {rule_id}",
            create_email_html(filtered_emails, current_folder),
            gr.update(),
            pending_rules, applied_rules, current_emails, sample_emails)






def handle_archive_rule(rule_id, mcp_client, pending_rules):
    """Handle archiving a rule via MCP"""
    if not rule_id or rule_id.strip() == "":
        return (
            create_interactive_rule_cards(pending_rules),
            create_preview_banner(get_preview_rules_count(pending_rules)),
            "❌ Error: No rule ID provided",
            pending_rules
        )
    
    try:
        # Archive via MCP
        response = mcp_client.archive_rule(rule_id)
        
        if response.get('success', False):
            # Re-fetch rules from backend
            updated_rules = mcp_client.get_rules()
            
            rule_name = response.get('rule', {}).get('name', 'Rule')
            status_msg = f"🗄️ Rule archived: {rule_name}"
            
            return (
                create_interactive_rule_cards(updated_rules),
                create_preview_banner(get_preview_rules_count(updated_rules)),
                status_msg,
                updated_rules
            )
        else:
            error_msg = response.get('error', 'Failed to archive')
            return (
                create_interactive_rule_cards(pending_rules),
                create_preview_banner(get_preview_rules_count(pending_rules)),
                f"❌ {error_msg}",
                pending_rules
            )
    except Exception as e:
        print(f"Error archiving rule: {e}")
        return (
            create_interactive_rule_cards(pending_rules),
            create_preview_banner(get_preview_rules_count(pending_rules)),
            f"❌ Connection error: {str(e)}",
            pending_rules
        )


def handle_reactivate_rule(rule_id, mcp_client, pending_rules):
    """Handle reactivating an archived rule via MCP"""
    if not rule_id or rule_id.strip() == "":
        return (
            create_interactive_rule_cards(pending_rules),
            create_preview_banner(get_preview_rules_count(pending_rules)),
            "❌ Error: No rule ID provided",
            pending_rules
        )
    
    try:
        # Reactivate via MCP
        response = mcp_client.reactivate_rule(rule_id)
        
        if response.get('success', False):
            # Re-fetch rules from backend
            updated_rules = mcp_client.get_rules()
            
            rule_name = response.get('rule', {}).get('name', 'Rule')
            status_msg = f"✅ Rule reactivated: {rule_name}"
            
            return (
                create_interactive_rule_cards(updated_rules),
                create_preview_banner(get_preview_rules_count(updated_rules)),
                status_msg,
                updated_rules
            )
        else:
            error_msg = response.get('error', 'Failed to reactivate')
            return (
                create_interactive_rule_cards(pending_rules),
                create_preview_banner(get_preview_rules_count(pending_rules)),
                f"❌ {error_msg}",
                pending_rules
            )
    except Exception as e:
        print(f"Error reactivating rule: {e}")
        return (
            create_interactive_rule_cards(pending_rules),
            create_preview_banner(get_preview_rules_count(pending_rules)),
            f"❌ Connection error: {str(e)}",
            pending_rules
        )


def exit_preview_mode(sort_option, search_query, pending_rules, current_emails, sample_emails):
    """Exit preview mode by setting all preview rules back to pending"""
    
    # Create new list of rules with updated status
    updated_rules = []
    for rule in pending_rules:
        if rule["status"] == "preview":
            # Create new rule object
            updated_rule = rule.copy()
            updated_rule["status"] = "pending"
            updated_rules.append(updated_rule)
        else:
            updated_rules.append(rule)
    
    current_emails = sample_emails
    
    new_folder_choices = get_folder_dropdown_choices(current_emails)
    new_folder_value = new_folder_choices[0] if new_folder_choices else "Inbox (0)"
    filtered_emails = filter_emails(new_folder_value, search_query, sort_option, current_emails)
    new_email_display = create_email_html(filtered_emails, new_folder_value)
    
    return (gr.update(choices=new_folder_choices, value=new_folder_value), 
            new_email_display, 
            create_interactive_rule_cards(updated_rules), 
            create_preview_banner(get_preview_rules_count(updated_rules)), 
            "🔄 Exited preview mode",
            updated_rules, current_emails)


def get_preview_rules_count(pending_rules):
    """Get the number of rules currently in preview status"""
    return len([rule for rule in pending_rules if rule["status"] == "preview"])


def create_app(modal_url: str) -> gr.Blocks:
    """
    Create the main Gradio app
    
    Args:
        modal_url: URL of the Modal backend server
        
    Returns:
        Gradio Blocks app
    """
    
    # Custom CSS for compact design with unified design system
    custom_css = """
    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap');
    @import url('https://cdn.jsdelivr.net/npm/toastify-js/src/toastify.min.css');
    
    /* Chat input alignment fix */
    .chat-input-row {
        display: flex;
        align-items: stretch;
    }
    
    .chat-input-row > div {
        display: flex !important;
        align-items: stretch !important;
        gap: 8px;
    }
    
    .chat-input-row textarea {
        min-height: 39.6px;
    }
    
    .chat-input-row button {
        height: auto !important;
        padding: 0.5rem 1rem !important;
        display: flex;
        align-items: center;
        justify-content: center;
        font-size: 1.2rem;
        min-height: 39.6px;
    }
    
    /* Use high-specificity selector to ensure .hide class works */
    gradio-app .hide {
        display: none !important;
    }
    
    /* Extra specificity for the main elements to prevent any override */
    #main_interface.hide,
    #todo_list_html.hide {
        display: none !important;
    }
    
    /* CSS Variables for Design System */
    :root {
        /* Colors */
        --color-primary: #007bff;
        --color-success: #28a745;
        --color-danger: #dc3545;
        --color-warning: #fd7e14;
        --color-info: #17a2b8;
        --color-dark: #212529;
        --color-gray: #6c757d;
        --color-gray-light: #e1e5e9;
        --color-gray-lighter: #f8f9fa;
        --color-white: #ffffff;
        
        /* Typography */
        --font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
        --font-size-xs: 0.75rem;    /* 12px */
        --font-size-sm: 0.875rem;   /* 14px */
        --font-size-base: 1rem;     /* 16px */
        --font-size-lg: 1.125rem;   /* 18px */
        --font-weight-normal: 400;
        --font-weight-medium: 500;
        --font-weight-semibold: 600;
        --font-weight-bold: 700;
        
        /* Spacing (4px base unit) */
        --spacing-1: 0.25rem;   /* 4px */
        --spacing-2: 0.5rem;    /* 8px */
        --spacing-3: 0.75rem;   /* 12px */
        --spacing-4: 1rem;      /* 16px */
        --spacing-5: 1.25rem;   /* 20px */
        --spacing-6: 1.5rem;    /* 24px */
        
        /* Borders */
        --border-width: 1px;
        --border-radius-sm: 4px;
        --border-radius: 6px;
        --border-radius-lg: 8px;
        
        /* Shadows */
        --shadow-sm: 0 1px 2px 0 rgba(0, 0, 0, 0.05);
        --shadow: 0 1px 3px 0 rgba(0, 0, 0, 0.1), 0 1px 2px 0 rgba(0, 0, 0, 0.06);
        --shadow-lg: 0 10px 15px -3px rgba(0, 0, 0, 0.1), 0 4px 6px -2px rgba(0, 0, 0, 0.05);
        
        /* Transitions */
        --transition-fast: 150ms ease-in-out;
        --transition-normal: 200ms ease-in-out;
        --transition-slow: 300ms ease-in-out;
    }
    
    /* Global Styles */
    * {
        font-family: var(--font-family) !important;
        box-sizing: border-box;
    }
    
    /* Layout - Modern Flexbox Approach */
    .main-container {
        height: calc(100vh - 100px);
        display: flex;
        flex-direction: column;
        overflow: hidden;
    }
    
    /* Main interface row should be flex container */
    .main-interface-row {
        display: flex;
        height: 100%;
        gap: var(--spacing-4);
        overflow: hidden;
    }
    
    /* Each column should be flex container */
    .main-interface-row > div[class*="gr-column"] {
        display: flex;
        flex-direction: column;
        height: 100%;
        overflow: hidden;
    }
    
    /* Typography */
    h3 {
        margin: 0;
        font-size: var(--font-size-base);
        font-weight: var(--font-weight-semibold);
    }
    
    /* Components */
    .compact-header {
        margin: var(--spacing-2) 0;
        display: flex;
        align-items: center;
        gap: var(--spacing-2);
        min-height: 32px; /* Ensure consistent height */
    }
    
    /* Fix header alignment across columns */
    .column-header-title {
        display: flex;
        align-items: center;
        height: 32px; /* Fixed height for all headers */
    }
    
    .compact-chat {
        flex-grow: 1;
        min-height: 300px;
        overflow-y: auto;
        border: var(--border-width) solid var(--color-gray-light);
        border-radius: var(--border-radius);
        background: var(--color-white);
    }
    
    .compact-input {
        font-size: var(--font-size-sm);
        padding: var(--spacing-2) var(--spacing-3);
        border: var(--border-width) solid var(--color-gray-light);
        border-radius: var(--border-radius-sm);
        transition: all var(--transition-fast);
    }
    
    .compact-input:focus {
        border-color: var(--color-primary);
        outline: none;
        box-shadow: 0 0 0 3px rgba(0, 123, 255, 0.1);
    }
    
    .compact-buttons {
        gap: var(--spacing-2);
        margin-top: var(--spacing-2);
    }
    
    .compact-buttons button,
    button[size="sm"] {
        padding: var(--spacing-2) var(--spacing-3);
        font-size: var(--font-size-sm);
        font-weight: var(--font-weight-medium);
        border-radius: var(--border-radius-sm);
        transition: all var(--transition-fast);
        cursor: pointer;
        position: relative;
    }
    
    .compact-buttons button:hover:not(:disabled) {
        transform: translateY(-1px);
        box-shadow: var(--shadow);
    }
    
    .compact-buttons button:active:not(:disabled) {
        transform: translateY(0);
        box-shadow: var(--shadow-sm);
    }
    
    
    .compact-rules {
        flex-grow: 1;
        overflow-y: auto;
        padding: var(--spacing-2);
    }
    
    .compact-rules::-webkit-scrollbar,
    .compact-emails::-webkit-scrollbar,
    .compact-chat::-webkit-scrollbar {
        width: 6px;
    }
    
    .compact-rules::-webkit-scrollbar-thumb,
    .compact-emails::-webkit-scrollbar-thumb,
    .compact-chat::-webkit-scrollbar-thumb {
        background: var(--color-gray-light);
        border-radius: 3px;
    }
    
    .compact-banner {
        margin-bottom: var(--spacing-2);
    }
    
    .compact-dropdown,
    .compact-search,
    .compact-sort {
        font-size: var(--font-size-sm);
        padding: var(--spacing-1) var(--spacing-2);
        border: var(--border-width) solid var(--color-gray-light);
        border-radius: var(--border-radius-sm);
        transition: border-color var(--transition-fast);
    }
    
    .compact-controls {
        gap: var(--spacing-2);
        margin: var(--spacing-2) 0;
    }
    
    .compact-emails {
        flex-grow: 1;
        overflow-y: auto;
        background: var(--color-white);
        border: var(--border-width) solid var(--color-gray-light);
        border-radius: var(--border-radius);
    }
    
    /* Gradio Overrides */
    .gr-form {
        gap: var(--spacing-2) !important;
    }
    
    .gr-box {
        padding: var(--spacing-3) !important;
        border-radius: var(--border-radius) !important;
    }
    
    .gr-padded {
        padding: var(--spacing-2) !important;
    }
    
    .gr-panel {
        padding: var(--spacing-2) !important;
    }
    
    .gr-button {
        font-family: var(--font-family) !important;
        font-weight: var(--font-weight-medium) !important;
        transition: all var(--transition-fast) !important;
        position: relative !important;
    }
    
    .gr-button:hover:not(:disabled) {
        transform: translateY(-1px);
        box-shadow: var(--shadow);
    }
    
    .gr-button:active:not(:disabled) {
        transform: translateY(0);
        box-shadow: var(--shadow-sm);
    }
    
    .gr-button[variant="primary"] {
        background: var(--color-primary) !important;
        border-color: var(--color-primary) !important;
    }
    
    .gr-button[variant="primary"]:hover:not(:disabled) {
        background: #0056b3 !important;
        border-color: #0056b3 !important;
    }
    
    .gr-button[variant="secondary"] {
        background: var(--color-white) !important;
        border: var(--border-width) solid var(--color-gray-light) !important;
        color: var(--color-dark) !important;
    }
    
    .gr-button[variant="secondary"]:hover:not(:disabled) {
        background: var(--color-gray-lighter) !important;
        border-color: var(--color-gray) !important;
    }
    
    /* Loading state for buttons */
    .gr-button:disabled {
        opacity: 0.6;
        cursor: not-allowed;
    }
    
    /* All interactive elements should have transitions */
    button, .gr-button, .email-row, .rule-card, input, select, textarea {
        transition: all var(--transition-fast);
    }
    
    /* Markdown styling */
    .gr-markdown {
        font-size: var(--font-size-sm);
        line-height: 1.5;
    }
    
    .gr-markdown h3 {
        font-size: var(--font-size-base);
        font-weight: var(--font-weight-semibold);
        margin: var(--spacing-2) 0;
    }
    
    /* Chatbot styling */
    .gr-chatbot {
        font-size: var(--font-size-sm);
    }
    
    .gr-chatbot .message {
        padding: var(--spacing-3);
        margin: var(--spacing-2) 0;
        border-radius: var(--border-radius);
    }
    
    /* Send button specific */
    button[min_width="45"] {
        min-width: 45px !important;
        width: 45px !important;
        padding: var(--spacing-2) !important;
        display: flex;
        align-items: center;
        justify-content: center;
    }
    
    /* Draft-specific styles */
    .draft-indicator {
        font-size: var(--font-size-sm);
        margin-right: var(--spacing-1);
        display: inline-flex;
        align-items: center;
    }
    
    .draft-preview {
        margin-top: var(--spacing-4);
        padding: var(--spacing-3);
        background-color: #e3f2fd;
        border-radius: var(--border-radius);
        border-left: 4px solid #2196f3;
        animation: fadeIn var(--transition-normal);
    }
    
    .draft-preview h4 {
        margin: 0 0 var(--spacing-2) 0;
        color: #1976d2;
        font-size: var(--font-size-sm);
        display: flex;
        align-items: center;
    }
    
    .draft-preview button {
        font-size: var(--font-size-xs);
        padding: 4px 12px;
        background: #2196f3;
        color: white;
        border: none;
        border-radius: var(--border-radius-sm);
        cursor: pointer;
        transition: all var(--transition-fast);
    }
    
    .draft-preview button:hover:not(:disabled) {
        background: #1976d2;
        transform: translateY(-1px);
        box-shadow: var(--shadow-sm);
    }
    
    .draft-preview button:active:not(:disabled) {
        transform: translateY(0);
        box-shadow: none;
    }
    
    @keyframes fadeIn {
        from { opacity: 0; transform: translateY(-10px); }
        to { opacity: 1; transform: translateY(0); }
    }
    
    /* New sleek two-row rule cards */
    .rule-card-v2 {
        border: 1px solid var(--color-gray-light);
        border-left-width: 4px;
        padding: var(--spacing-2) var(--spacing-3);
        margin-bottom: var(--spacing-2);
        background: var(--color-white);
        transition: all var(--transition-fast);
        border-radius: var(--border-radius-sm);
        cursor: pointer;
    }
    
    .rule-card-v2:hover {
        box-shadow: var(--shadow);
        transform: translateY(-1px);
    }
    
    /* Expand arrow */
    .rule-expand-arrow {
        display: inline-block;
        width: 16px;
        color: var(--color-gray);
        font-size: var(--font-size-xs);
        transition: transform var(--transition-fast);
        margin-right: var(--spacing-2);
    }
    
    .rule-expand-arrow.expanded {
        transform: rotate(90deg);
    }
    
    /* Rule details section */
    .rule-details {
        padding-top: var(--spacing-2);
        border-top: 1px solid var(--color-gray-light);
        margin-top: var(--spacing-2);
        font-size: var(--font-size-xs);
        color: var(--color-dark);
    }
    
    .rule-details strong {
        color: var(--color-gray);
        font-weight: var(--font-weight-semibold);
    }
    
    .rule-details ul {
        list-style-type: disc;
        margin: var(--spacing-1) 0;
    }
    
    .rule-details li {
        color: var(--color-dark);
        line-height: 1.4;
    }
    
    /* Status-specific left border colors */
    .rule-status-pending { border-left-color: var(--color-gray); }
    .rule-status-preview { border-left-color: var(--color-warning); }
    .rule-status-accepted { border-left-color: var(--color-success); }
    
    .rule-row-1, .rule-row-2 {
        display: flex;
        align-items: center;
        justify-content: space-between;
    }
    
    .rule-row-1 {
        margin-bottom: var(--spacing-2);
    }
    
    .rule-name {
        font-weight: var(--font-weight-semibold);
        font-size: var(--font-size-sm);
        color: var(--color-dark);
    }
    
    .rule-actions {
        display: flex;
        gap: var(--spacing-2);
    }
    
    /* Button specific styles */
    .rule-btn {
        padding: 6px 12px;
        font-size: var(--font-size-xs);
        font-weight: var(--font-weight-medium);
        border-radius: var(--border-radius-sm);
        transition: all var(--transition-fast);
        cursor: pointer;
        border: 1px solid transparent;
    }
    
    .rule-btn:hover {
        transform: translateY(-1px);
        box-shadow: var(--shadow-sm);
    }
    
    .rule-btn.accept-btn {
        background-color: var(--color-success);
        color: white;
        border-color: var(--color-success);
    }
    
    .rule-btn.accept-btn:hover {
        background-color: #218838;
    }
    
    .rule-btn.reject-btn {
        background-color: var(--color-danger);
        color: white;
        border-color: var(--color-danger);
    }
    
    .rule-btn.reject-btn:hover {
        background-color: #c82333;
    }
    
    .rule-btn.run-btn, .rule-btn.preview-btn {
        background-color: var(--color-white);
        color: var(--color-primary);
        border: 1px solid var(--color-primary);
    }
    
    .rule-btn.run-btn:hover, .rule-btn.preview-btn:hover {
        background-color: var(--color-primary);
        color: white;
    }
    
    .rule-btn.archive-btn {
        background-color: transparent;
        color: var(--color-gray);
        border: 1px solid var(--color-gray-light);
    }
    
    .rule-btn.archive-btn:hover {
        background-color: var(--color-gray-lighter);
        border-color: var(--color-gray);
    }
    
    /* Demo Checklist Styles */
    .demo-checklist {
        position: fixed;
        top: 80px;
        right: 20px;
        width: 320px;
        background: var(--color-white);
        border: 2px solid var(--color-primary);
        border-radius: var(--border-radius-lg);
        box-shadow: 0 10px 30px rgba(0, 0, 0, 0.15);
        z-index: 1000;
        display: none;
        flex-direction: column;
        transition: all var(--transition-normal);
    }
    
    .demo-checklist.show {
        display: flex;
    }
    
    .demo-checklist.initial-show {
        animation: slideInRight 0.5s ease-out forwards;
    }
    
    .demo-checklist.minimized {
        height: auto;
    }
    
    .demo-checklist.minimized .checklist-body {
        display: none;
    }
    
    @keyframes slideInRight {
        from {
            transform: translateX(400px);
            opacity: 0;
        }
        to {
            transform: translateX(0);
            opacity: 1;
        }
    }
    
    @keyframes pulse {
        0%, 100% { transform: scale(1); }
        50% { transform: scale(1.05); }
    }
    
    .demo-checklist.pulse-animation {
        animation: pulse 0.5s ease-out 3;
    }
    
    .checklist-header {
        padding: var(--spacing-3) var(--spacing-4);
        background: linear-gradient(135deg, var(--color-primary), #0056b3);
        color: var(--color-white);
        cursor: move;
        border-radius: var(--border-radius) var(--border-radius) 0 0;
        display: flex;
        justify-content: space-between;
        align-items: center;
        font-weight: var(--font-weight-bold);
        font-size: var(--font-size-base);
        user-select: none;
    }
    
    .checklist-header:hover {
        background: linear-gradient(135deg, #0056b3, var(--color-primary));
    }
    
    .checklist-header:active {
        cursor: grabbing;
    }
    
    .minimize-btn {
        background: rgba(255, 255, 255, 0.2);
        border: none;
        color: var(--color-white);
        padding: 4px 8px;
        border-radius: var(--border-radius-sm);
        cursor: pointer;
        font-size: var(--font-size-sm);
        transition: all var(--transition-fast);
    }
    
    .minimize-btn:hover {
        background: rgba(255, 255, 255, 0.3);
        transform: scale(1.1);
    }
    
    .checklist-body {
        padding: var(--spacing-3);
        overflow-y: auto;
        flex: 1;
        max-height: 350px;
    }
    
    .progress-section {
        margin-bottom: var(--spacing-3);
        padding: var(--spacing-2);
        background: var(--color-gray-lighter);
        border-radius: var(--border-radius-sm);
    }
    
    .progress-label {
        font-size: var(--font-size-sm);
        font-weight: var(--font-weight-medium);
        color: var(--color-dark);
        margin-bottom: var(--spacing-1);
    }
    
    .progress-bar {
        width: 100%;
        height: 8px;
        background: var(--color-gray-light);
        border-radius: 4px;
        overflow: hidden;
    }
    
    .progress-fill {
        height: 100%;
        background: linear-gradient(90deg, var(--color-success), #22c55e);
        width: 0%;
        transition: width var(--transition-normal);
    }
    
    .checklist-item {
        padding: var(--spacing-2) var(--spacing-1);
        border-radius: var(--border-radius-sm);
        display: flex;
        align-items: center;
        gap: var(--spacing-2);
        transition: all var(--transition-fast);
        cursor: pointer;
    }
    
    .checklist-item:hover {
        background: var(--color-gray-lighter);
        transform: translateX(4px);
    }
    
    .checklist-checkbox {
        width: 18px;
        height: 18px;
        cursor: pointer;
        accent-color: var(--color-success);
    }
    
    .checklist-text {
        flex: 1;
        font-size: var(--font-size-sm);
        transition: all var(--transition-fast);
        user-select: none;
    }
    
    .checklist-item.completed .checklist-text {
        text-decoration: line-through;
        color: var(--color-gray);
        opacity: 0.7;
    }
    
    .checklist-item.completed {
        background: rgba(40, 167, 69, 0.1);
    }
    
    /* Animation for completion */
    @keyframes checkComplete {
        0% { transform: scale(1) rotate(0deg); }
        50% { transform: scale(1.2) rotate(5deg); }
        100% { transform: scale(1) rotate(0deg); }
    }
    
    .checklist-item.completed .checklist-checkbox {
        animation: checkComplete 0.3s ease-out;
    }
    
    /* Header progress indicator */
    .header-progress {
        display: none;
        font-size: 0.9rem;
        color: rgba(255, 255, 255, 0.9);
        margin: 0 10px;
    }
    
    .demo-checklist.minimized .header-progress {
        display: inline;
    }
    
    /* Celebration Overlay Styles */
    .celebration-overlay {
        position: absolute;
        top: 0;
        left: 0;
        width: 100%;
        height: 100%;
        background-color: rgba(0, 0, 0, 0.5);
        display: flex;
        justify-content: center;
        align-items: center;
        opacity: 0;
        pointer-events: none;
        transition: opacity 0.3s ease-in-out;
        border-radius: var(--border-radius-lg);
        z-index: 10;
    }
    
    .celebration-overlay:not(.hidden) {
        opacity: 1;
        pointer-events: auto;
    }
    
    .celebration-content {
        background: white;
        padding: 2rem;
        border-radius: 12px;
        box-shadow: 0 10px 30px rgba(0,0,0,0.2);
        text-align: center;
        transform: scale(0.8);
        transition: transform 0.3s cubic-bezier(0.175, 0.885, 0.32, 1.275);
        position: relative;
    }
    
    .celebration-overlay:not(.hidden) .celebration-content {
        transform: scale(1);
    }
    
    .celebration-close {
        position: absolute;
        top: 10px;
        right: 10px;
        background: none;
        border: none;
        font-size: 24px;
        cursor: pointer;
        color: #666;
        transition: color 0.2s;
    }
    
    .celebration-close:hover {
        color: #333;
    }
    
    .celebration-heading {
        margin: 0 0 1rem 0;
        font-size: 1.5rem;
        color: var(--color-primary);
    }
    
    .celebration-message {
        margin: 0;
        font-size: 1rem;
        color: #555;
    }
    """
    
    # Create the Gradio interface
    with gr.Blocks(title="Email Rule Agent", theme=gr.themes.Soft(), css=custom_css) as demo:
        # Initialize state variables
        session_id = gr.State(value=str(uuid.uuid4()))
        mcp_client_state = gr.State()
        
        # Email state
        sample_emails_state = gr.State([])
        preview_emails_state = gr.State([])
        current_emails_state = gr.State([])
        expanded_email_id_state = gr.State(None)
        
        # Rule state
        pending_rules_state = gr.State([])
        applied_rules_state = gr.State([])
        rule_counter_state = gr.State(0)
        
        gr.Markdown("# 📧 Email Rule Agent")
        gr.Markdown("*Intelligent email management powered by AI*")
        
        # Add login choice
        with gr.Row(visible=True, elem_id="login_row") as login_row:
            with gr.Column():
                gr.Markdown("## Welcome! Choose how to get started:")
                with gr.Row():
                    demo_btn = gr.Button("🎮 Try with Demo Data", variant="primary", scale=1)
                    gmail_btn = gr.Button("📧 Login with Gmail *Work in Progress*", variant="secondary", scale=1)
        
        # Loading indicator (initially hidden)
        with gr.Column(visible=False, elem_id="loading_indicator") as loading_indicator:
            gr.HTML("""
                <div style="text-align: center; padding: 50px;">
                    <div style="display: inline-block; width: 50px; height: 50px; border: 5px solid #f3f3f3; border-radius: 50%; border-top: 5px solid #007bff; animation: spin 1s linear infinite;"></div>
                    <h3 style="margin-top: 20px; color: #666;">Initializing Email Rule Agent...</h3>
                    <p style="color: #999;">Connecting to backend and loading your emails...</p>
                </div>
                <style>
                    @keyframes spin {
                        0% { transform: rotate(0deg); }
                        100% { transform: rotate(360deg); }
                    }
                </style>
            """)
        
        # Main interface (initially hidden)
        with gr.Column(visible=False, elem_id="main_interface") as main_interface:
            # Preview banner above all columns
            preview_banner = gr.HTML(create_preview_banner(0), elem_classes="compact-banner")
            
            # 3 column layout
            with gr.Row(elem_classes="main-interface-row"):
                # Left column - Chat Interface (35%)
                with gr.Column(scale=35, min_width=300):
                    gr.Markdown("### 💬 Chat", elem_classes="compact-header column-header-title")
                    chatbot = gr.Chatbot(
                        value=[],
                        type="messages",
                        height="40vh",
                        render_markdown=True,
                        sanitize_html=False,  # Allow our custom HTML for tool messages
                        elem_classes="compact-chat",
                        label=None,
                        show_label=False
                    )
                    
                    with gr.Row(elem_classes="chat-input-row"):
                        msg_input = gr.Textbox(
                            placeholder="Describe how you want to organize your emails...",
                            container=False,
                            elem_classes="compact-input",
                            scale=8
                        )
                        send_btn = gr.Button("❯", scale=1, variant="primary", min_width=45)
                    
                    # Quick action buttons with text
                    with gr.Row(elem_classes="compact-buttons"):
                        analyze_btn = gr.Button("🔍 Analyze Inbox *WIP*", scale=1, size="sm")
                        draft_btn = gr.Button("📝 Draft Replies *WIP*", scale=1, size="sm")
                
                # Middle column - Rules (25%)
                with gr.Column(scale=25, min_width=250):
                    gr.Markdown("### 📋 Rules", elem_classes="compact-header column-header-title")
                    
                    # Hidden status message (we'll show messages inline or as toasts)
                    status_msg = gr.Textbox(
                        value="",
                        visible=False,
                        elem_id="status_msg"
                    )
                
                    # Hidden components for JavaScript interaction
                    with gr.Row(visible=False):
                        preview_btn = gr.Button("Preview", elem_id="hidden_preview_btn")
                        accept_btn = gr.Button("Accept", elem_id="hidden_accept_btn")
                        reject_btn = gr.Button("Reject", elem_id="hidden_reject_btn")
                        exit_preview_btn = gr.Button("Exit Preview", elem_id="hidden_exit_preview_btn")
                        run_rule_btn = gr.Button("Run Rule", elem_id="hidden_run_rule_btn")
                        archive_rule_btn = gr.Button("Archive Rule", elem_id="hidden_archive_rule_btn")
                        reactivate_rule_btn = gr.Button("Reactivate Rule", elem_id="hidden_reactivate_rule_btn")
                        current_rule_id = gr.Textbox(visible=False, elem_id="current_rule_id")
                        
                        # For email expansion
                        expand_email_btn = gr.Button("Expand", elem_id="hidden_expand_email_btn")
                        current_email_id_input = gr.Textbox(visible=False, elem_id="current_email_id_input")

                    # Rule cards display with scroll
                    rule_cards = gr.HTML(
                        create_interactive_rule_cards([]),
                        elem_classes="compact-rules"
                    )
                
                # Right column - Email Display (40%)
                with gr.Column(scale=40, min_width=350):
                    # Email header now at the top like other columns
                    gr.Markdown("### 📧 Emails", elem_classes="compact-header column-header-title")
                    
                    # Folder dropdown (separate row)
                    with gr.Row(elem_classes="compact-controls"):
                        folder_dropdown = gr.Dropdown(
                            choices=["Inbox (0)"],
                            value="Inbox (0)",
                            container=False,
                            scale=1,
                            elem_classes="compact-dropdown"
                        )
                    
                    # Search and sort
                    with gr.Row(elem_classes="compact-controls"):
                        search_box = gr.Textbox(
                            placeholder="Search...",
                            container=False,
                            scale=3,
                            elem_classes="compact-search"
                        )
                        sort_dropdown = gr.Dropdown(
                            choices=["Newest", "Oldest"],
                            value="Newest",
                            container=False,
                            scale=1,
                            elem_classes="compact-sort"
                        )
                    
                    # Email list with scroll
                    email_display = gr.HTML(
                        value=create_email_html([]),
                        elem_classes="compact-emails"
                    )
        
        # Add Demo Checklist HTML (initially hidden)
        todo_list_html = gr.HTML("""
        <div class="demo-checklist" id="demo-checklist">
            <div class="checklist-header" onmousedown="startDrag(event)">
                <span>🚀 Demo Checklist</span>
                <span class="header-progress" id="header-progress"></span>
                <button class="minimize-btn" id="minimize-btn" onclick="toggleChecklist(event)">▼</button>
            </div>
            <div class="checklist-body">
                <div class="progress-section">
                    <div class="progress-label">Progress: <span id="progress-text">0/5 completed</span></div>
                    <div class="progress-bar">
                        <div class="progress-fill" id="progress-fill"></div>
                    </div>
                </div>
                <div class="checklist-items">
                    <div class="checklist-item" data-action="manual-rule">
                        <input type="checkbox" class="checklist-checkbox" onchange="updateProgress()">
                        <span class="checklist-text">🚗 Create rule: "Move Uber receipts to travel folder"</span>
                    </div>
                    <div class="checklist-item" data-action="preview">
                        <input type="checkbox" class="checklist-checkbox" onchange="updateProgress()">
                        <span class="checklist-text">👁️ Preview the rule to see affected emails</span>
                    </div>
                    <div class="checklist-item" data-action="accept">
                        <input type="checkbox" class="checklist-checkbox" onchange="updateProgress()">
                        <span class="checklist-text">✅ Accept the rule and organize emails</span>
                    </div>
                </div>
            </div>
            <div id="celebration-overlay" class="celebration-overlay hidden">
                <div class="celebration-content">
                    <button class="celebration-close" onclick="dismissCelebration()">×</button>
                    <h3 class="celebration-heading">🎉 All Done! 🎉</h3>
                    <p class="celebration-message">Thank you for demoing our hackathon app!</p>
                </div>
            </div>
        </div>
        """, visible=False, elem_id="todo_list_html")
        
        # Initialize data on load
        def initialize_app(session_id):
            """Initialize app with data from MCP backend"""
            # Initialize MCP client with session
            mcp_client = MCPClient(modal_url, session_token=session_id)
            
            # Fetch initial data from backend
            email_response = mcp_client.list_emails(folder='inbox', page_size=50)
            initial_rules = mcp_client.get_rules()
            
            # Extract data from responses
            initial_emails = email_response.get('emails', [])
            
            # For preview functionality, fetch a different set
            preview_response = mcp_client.list_emails(folder='inbox', page_size=20)
            preview_emails = preview_response.get('emails', [])
            
            # Initialize rule counter based on existing rules
            rule_counter = len(initial_rules)
            
            # Create welcome message
            welcome_chat = [{
                "role": "assistant",
                "content": """👋 Hi! I'm your Email Organization Assistant. I can help you:

• 📋 **Create smart filters** to automatically organize emails  
• 🏷️ **Set up labels** for different types of emails
• 📁 **Move emails** to folders based on sender, subject, or content

Try clicking "Analyze" below or tell me how you'd like to organize your emails!"""
            }]
            
            # Update UI components
            folder_choices = get_folder_dropdown_choices(initial_emails)
            initial_folder = folder_choices[0] if folder_choices else "Inbox (0)"
            initial_email_html = create_email_html(initial_emails)
            initial_rule_cards = create_interactive_rule_cards(initial_rules)
            initial_preview_banner = create_preview_banner(get_preview_rules_count(initial_rules))
            
            return (
                mcp_client,  # mcp_client_state
                initial_emails,  # sample_emails_state
                preview_emails,  # preview_emails_state
                initial_emails,  # current_emails_state
                initial_rules,  # pending_rules_state
                [],  # applied_rules_state
                rule_counter,  # rule_counter_state
                gr.update(choices=folder_choices, value=initial_folder),  # folder_dropdown
                initial_email_html,  # email_display
                initial_rule_cards,  # rule_cards
                initial_preview_banner,  # preview_banner
                welcome_chat,  # chatbot initial value
                None # expanded_email_id_state
            )
        
        # Event handlers
        def start_demo_mode(session_id):
            try:
                # Initialize and set session to demo mode
                results = initialize_app(session_id)
                mcp_client = results[0]
                mcp_client.set_mode('mock')
                return (
                    *results,
                    gr.update(visible=False),  # login_row
                    gr.update(visible=False),  # loading_indicator
                    gr.update(visible=True),   # main_interface
                    gr.update(visible=True)    # todo_list_html
                )
            except Exception as e:
                print(f"Error initializing app: {e}")
                error_chat = [{
                    "role": "assistant",
                    "content": f"❌ **Failed to initialize the app**\n\nError: {str(e)}\n\nPlease try refreshing the page or running in local mode with `python app.py --local`"
                }]
                # Return minimal state with error message
                return (
                    MCPClient('local://mock'),  # mcp_client_state
                    [],  # sample_emails_state
                    [],  # preview_emails_state
                    [],  # current_emails_state
                    [],  # pending_rules_state
                    [],  # applied_rules_state
                    0,   # rule_counter_state
                    gr.update(choices=["Inbox (0)"], value="Inbox (0)"),  # folder_dropdown
                    create_email_html([]),  # email_display
                    create_interactive_rule_cards([]),  # rule_cards
                    create_preview_banner(0),  # preview_banner
                    error_chat,  # chatbot
                    None,  # expanded_email_id_state
                    gr.update(visible=True),   # login_row (show again)
                    gr.update(visible=False),  # loading_indicator
                    gr.update(visible=False),  # main_interface
                    gr.update(visible=False)   # todo_list_html
                )
        
        def start_gmail_mode(session_id):
            try:
                # Initialize and get Gmail OAuth URL
                results = initialize_app(session_id)
                mcp_client = results[0]
                auth_url = mcp_client.get_gmail_auth_url()
                if auth_url:
                    # In a real app, we'd redirect to this URL
                    print(f"Gmail auth URL: {auth_url}")
                return (
                    *results,
                    gr.update(visible=False),  # login_row
                    gr.update(visible=False),  # loading_indicator
                    gr.update(visible=True),   # main_interface
                    gr.update(visible=True)    # todo_list_html
                )
            except Exception as e:
                print(f"Error initializing Gmail mode: {e}")
                # Return same error state as demo mode
                error_chat = [{
                    "role": "assistant",
                    "content": f"❌ **Failed to initialize Gmail mode**\n\nError: {str(e)}\n\nPlease try the demo mode or check your connection."
                }]
                return (
                    MCPClient('local://mock'),  # mcp_client_state
                    [],  # sample_emails_state
                    [],  # preview_emails_state
                    [],  # current_emails_state
                    [],  # pending_rules_state
                    [],  # applied_rules_state
                    0,   # rule_counter_state
                    gr.update(choices=["Inbox (0)"], value="Inbox (0)"),  # folder_dropdown
                    create_email_html([]),  # email_display
                    create_interactive_rule_cards([]),  # rule_cards
                    create_preview_banner(0),  # preview_banner
                    error_chat,  # chatbot
                    None,  # expanded_email_id_state
                    gr.update(visible=True),   # login_row (show again)
                    gr.update(visible=False),  # loading_indicator
                    gr.update(visible=False),  # main_interface
                    gr.update(visible=False)   # todo_list_html
                )
        
        # Show loading function
        def show_loading():
            return (
                gr.update(visible=False),  # login_row
                gr.update(visible=True)    # loading_indicator
            )
        
        demo_btn.click(
            show_loading,
            inputs=None,
            outputs=[login_row, loading_indicator],
            queue=False
        ).then(
            start_demo_mode,
            inputs=[session_id],
            outputs=[
                mcp_client_state, sample_emails_state, preview_emails_state, current_emails_state,
                pending_rules_state, applied_rules_state, rule_counter_state,
                folder_dropdown, email_display, rule_cards, preview_banner, chatbot,
                expanded_email_id_state,
                login_row, loading_indicator, main_interface, todo_list_html
            ]
        ).then(
            None,
            None,
            None,
            js="() => { setTimeout(() => window.showDemoChecklist(), 100); }"
        )
        
        gmail_btn.click(
            show_loading,
            inputs=None,
            outputs=[login_row, loading_indicator],
            queue=False
        ).then(
            start_gmail_mode,
            inputs=[session_id],
            outputs=[
                mcp_client_state, sample_emails_state, preview_emails_state, current_emails_state,
                pending_rules_state, applied_rules_state, rule_counter_state,
                folder_dropdown, email_display, rule_cards, preview_banner, chatbot,
                expanded_email_id_state,
                login_row, loading_indicator, main_interface, todo_list_html
            ]
        )
        
        # Use the streaming chat handler
        
        send_btn.click(
            process_chat_streaming,
            inputs=[msg_input, chatbot, mcp_client_state, current_emails_state,
                   pending_rules_state, applied_rules_state, rule_counter_state],
            outputs=[chatbot, msg_input, rule_cards, status_msg,
                    pending_rules_state, rule_counter_state],
            show_progress="full"  # Show progress for streaming
        )
        
        msg_input.submit(
            process_chat_streaming,
            inputs=[msg_input, chatbot, mcp_client_state, current_emails_state,
                   pending_rules_state, applied_rules_state, rule_counter_state],
            outputs=[chatbot, msg_input, rule_cards, status_msg,
                    pending_rules_state, rule_counter_state],
            show_progress="full"
        )
        
        # Quick action handlers with streaming
        def analyze_inbox(hist, mcp, emails, pend, appl, cnt):
            yield from process_chat_streaming(
                "Analyze my inbox and suggest organization rules", 
                hist, mcp, emails, pend, appl, cnt
            )
        
        analyze_btn.click(
            analyze_inbox,
            inputs=[chatbot, mcp_client_state, current_emails_state,
                   pending_rules_state, applied_rules_state, rule_counter_state],
            outputs=[chatbot, msg_input, rule_cards, status_msg,
                    pending_rules_state, rule_counter_state],
            show_progress="full"
        )
        
        def draft_replies(hist, mcp, emails, pend, appl, cnt):
            yield from process_chat_streaming(
                "Draft polite replies to all meeting invitations saying I'll check my calendar and get back to them", 
                hist, mcp, emails, pend, appl, cnt
            )
        
        draft_btn.click(
            draft_replies,
            inputs=[chatbot, mcp_client_state, current_emails_state,
                   pending_rules_state, applied_rules_state, rule_counter_state],
            outputs=[chatbot, msg_input, rule_cards, status_msg,
                    pending_rules_state, rule_counter_state],
            show_progress="full"
        )
        
        # Email display handlers
        def update_emails(folder, sort_option, search_query, current_emails):
            filtered_emails = filter_emails(folder, search_query, sort_option, current_emails)
            return create_email_html(filtered_emails, folder)
        
        folder_dropdown.change(
            update_emails,
            inputs=[folder_dropdown, sort_dropdown, search_box, current_emails_state],
            outputs=[email_display]
        )
        
        sort_dropdown.change(
            lambda folder, sort, search, emails: update_emails(
                folder, 
                "Newest First" if sort == "Newest" else "Oldest First",
                search,
                emails
            ),
            inputs=[folder_dropdown, sort_dropdown, search_box, current_emails_state],
            outputs=[email_display]
        )
        
        search_box.change(
            update_emails,
            inputs=[folder_dropdown, sort_dropdown, search_box, current_emails_state],
            outputs=[email_display]
        )
        
        # Rule action handlers
        preview_btn.click(
            handle_preview_rule,
            inputs=[current_rule_id, sort_dropdown, search_box, mcp_client_state,
                   pending_rules_state, current_emails_state, sample_emails_state, preview_emails_state],
            outputs=[folder_dropdown, email_display, rule_cards, preview_banner, status_msg,
                    pending_rules_state, current_emails_state]
        )
        
        accept_btn.click(
            lambda rule_id, folder, sort, search, mcp, pend, appl, curr, samp: handle_accept_rule(
                rule_id, mcp, folder, sort, search, pend, appl, curr, samp),
            inputs=[current_rule_id, folder_dropdown, sort_dropdown, search_box, mcp_client_state,
                   pending_rules_state, applied_rules_state, current_emails_state, sample_emails_state],
            outputs=[rule_cards, preview_banner, status_msg, email_display, folder_dropdown,
                    pending_rules_state, applied_rules_state, current_emails_state, sample_emails_state]
        )
        
        reject_btn.click(
            handle_archive_rule,
            inputs=[current_rule_id, mcp_client_state, pending_rules_state],
            outputs=[rule_cards, preview_banner, status_msg, pending_rules_state]
        )
        
        exit_preview_btn.click(
            lambda sort, search, pend, curr, samp: exit_preview_mode(sort, search, pend, curr, samp),
            inputs=[sort_dropdown, search_box, pending_rules_state, current_emails_state, sample_emails_state],
            outputs=[folder_dropdown, email_display, rule_cards, preview_banner, status_msg,
                    pending_rules_state, current_emails_state]
        )
        
        run_rule_btn.click(
            lambda rule_id, folder, sort, search, mcp, pend, appl, curr, samp: handle_run_rule(
                rule_id, mcp, folder, sort, search, pend, appl, curr, samp),
            inputs=[current_rule_id, folder_dropdown, sort_dropdown, search_box, mcp_client_state,
                   pending_rules_state, applied_rules_state, current_emails_state, sample_emails_state],
            outputs=[rule_cards, preview_banner, status_msg, email_display, folder_dropdown,
                    pending_rules_state, applied_rules_state, current_emails_state, sample_emails_state]
        )
        
        archive_rule_btn.click(
            handle_archive_rule,
            inputs=[current_rule_id, mcp_client_state, pending_rules_state],
            outputs=[rule_cards, preview_banner, status_msg, pending_rules_state]
        )
        
        reactivate_rule_btn.click(
            handle_reactivate_rule,
            inputs=[current_rule_id, mcp_client_state, pending_rules_state],
            outputs=[rule_cards, preview_banner, status_msg, pending_rules_state]
        )
        
        # Email expansion handler
        def toggle_email_expansion_handler(email_id, folder, sort, search, current_emails, current_expanded_id):
            new_expanded_id = email_id if email_id != current_expanded_id else None
            filtered_emails = filter_emails(folder, search, sort, current_emails)
            html = create_email_html(filtered_emails, folder, expanded_email_id=new_expanded_id)
            return html, new_expanded_id

        expand_email_btn.click(
            toggle_email_expansion_handler,
            inputs=[current_email_id_input, folder_dropdown, sort_dropdown, search_box, current_emails_state, expanded_email_id_state],
            outputs=[email_display, expanded_email_id_state]
        )
        
        # Load JavaScript handlers
        demo.load(
            None,
            inputs=None,
            outputs=None,
            js=get_javascript_handlers()
        )
    
    return demo


def launch_app(demo: gr.Blocks):
    """Launch the Gradio app"""
    demo.launch()


def get_javascript_handlers():
    """Get JavaScript code for handling UI interactions"""
    return """
    () => {
        // Load Toastify.js dynamically
        const script = document.createElement('script');
        script.src = 'https://cdn.jsdelivr.net/npm/toastify-js';
        document.head.appendChild(script);
        
        // Toast notification function
        window.showToast = function(message, type = 'info') {
            if (!window.Toastify) {
                console.log('Toast:', message);
                return;
            }
            
            let backgroundColor = "linear-gradient(to right, #00b09b, #96c93d)"; // Default success
            
            // Determine color based on message content or type
            if (message.startsWith("❌") || message.toLowerCase().includes("error") || type === 'error') {
                backgroundColor = "linear-gradient(to right, #ff5f6d, #ffc371)"; // Error
            } else if (message.startsWith("⚠️") || message.startsWith("👁️") || type === 'warning') {
                backgroundColor = "linear-gradient(to right, #ffd166, #f77f00)"; // Warning
            } else if (message.startsWith("✅") || message.startsWith("✓") || type === 'success') {
                backgroundColor = "linear-gradient(to right, #00b09b, #96c93d)"; // Success
            } else if (message.startsWith("🗄️") || message.startsWith("🔄") || type === 'info') {
                backgroundColor = "linear-gradient(to right, #667eea, #764ba2)"; // Info
            }
            
            Toastify({
                text: message,
                duration: 3000,
                close: true,
                gravity: "top",
                position: "right",
                stopOnFocus: true,
                style: {
                    background: backgroundColor,
                },
                onClick: function(){} // Callback after click
            }).showToast();
        };
        
        // Set up observer for status_msg changes
        setTimeout(() => {
            const statusContainer = document.getElementById('status_msg');
            if (statusContainer) {
                const statusTextarea = statusContainer.querySelector('textarea');
                if (statusTextarea) {
                    // Previous value to detect changes
                    let previousValue = statusTextarea.value;
                    
                    // List of status messages to ignore (agent typing/processing messages)
                    const ignoredStatuses = [
                        'Processing...', 
                        'Typing...', 
                        'Analyzing emails...', 
                        'Creating rules...',
                        'Ready...',
                        'Complete',
                        'API key missing'
                    ];
                    
                    // Track if we've shown a toast recently to prevent duplicates
                    let lastToastTime = 0;
                    const TOAST_DEBOUNCE = 500; // 500ms between toasts
                    
                    // Function to handle status changes
                    const handleStatusChange = () => {
                        const currentValue = statusTextarea.value;
                        const now = Date.now();
                        
                        if (currentValue && currentValue.trim() && currentValue !== previousValue) {
                            // Check if we've shown a toast recently
                            if (now - lastToastTime < TOAST_DEBOUNCE) {
                                return;
                            }
                            
                            // Don't show toast for agent typing/processing messages
                            const isIgnored = ignoredStatuses.some(status => 
                                currentValue === status || currentValue.startsWith(status)
                            );
                            
                            // Only show toast for actual status updates (errors, rule actions, etc.)
                            if (!isIgnored && !currentValue.includes('Found') && !currentValue.includes('rules')) {
                                window.showToast(currentValue);
                                lastToastTime = now;
                            }
                            
                            previousValue = currentValue;
                            // Clear after showing toast
                            setTimeout(() => {
                                statusTextarea.value = '';
                                previousValue = '';
                            }, 200);
                        }
                    };
                    
                    // Listen for input events
                    statusTextarea.addEventListener('input', handleStatusChange);
                    
                    // Also check periodically as backup (less frequent)
                    setInterval(handleStatusChange, 300);
                    
                    console.log('Toast notification system initialized');
                } else {
                    console.error('Status textarea not found');
                }
            } else {
                console.error('Status container not found');
            }
        }, 1000);
        
        // Expose email handler to global scope
        window.handleEmailAction = (emailId) => {
            console.log(`handleEmailAction called with emailId: ${emailId}`);
            
            const container = document.getElementById('current_email_id_input');
            let emailIdInput = null;
            if (container) {
                emailIdInput = container.querySelector('input, textarea');
            }

            if (emailIdInput) {
                emailIdInput.value = emailId;
                emailIdInput.dispatchEvent(new Event('input', { bubbles: true }));
                emailIdInput.dispatchEvent(new Event('change', { bubbles: true }));
                
                setTimeout(() => {
                    const targetButton = document.getElementById('hidden_expand_email_btn');
                    if (targetButton) {
                        targetButton.click();
                    } else {
                        console.error('Hidden expand email button not found');
                    }
                }, 100);
            } else {
                console.error('Email ID input field not found');
            }
        };

        // Set up global click handler once
        console.log('Setting up global click handler');
        document.addEventListener('click', globalClickHandler);
        
        function globalClickHandler(e) {
            // Use a single selector for all action buttons
            const actionBtn = e.target.closest('.preview-btn, .accept-btn, .reject-btn, .run-btn, .archive-btn, .reactivate-btn');
            
            if (actionBtn) {
                e.preventDefault();
                e.stopPropagation(); // Stop propagation AFTER we handle it
                
                // Extract action from class name (skip 'rule-btn')
                const action = Array.from(actionBtn.classList)
                    .find(c => c.endsWith('-btn') && c !== 'rule-btn')
                    .replace('-btn', '');
                
                const ruleId = actionBtn.getAttribute('data-rule-id');
                console.log(`${action} button clicked for rule: ${ruleId}`);
                
                handleRuleAction(action, ruleId);
            } else if (e.target.closest('#exit-preview-btn, .exit-preview-card-btn')) {
                e.preventDefault();
                e.stopPropagation();
                handleExitPreview();
            }
        }
        
        function handleExitPreview() {
            console.log('Exit preview button clicked');
            const exitButton = document.getElementById('hidden_exit_preview_btn');
            if (exitButton) {
                exitButton.click();
            } else {
                console.error('Hidden exit preview button not found');
            }
        }
        
        function handleRuleAction(action, ruleId) {
            console.log(`handleRuleAction called with action: ${action}, ruleId: ${ruleId}`);
            
            if (!ruleId) {
                console.error('No ruleId provided');
                window.showToast('❌ No rule ID provided');
                return;
            }
            
            try {
                // Find the hidden textbox
                const container = document.getElementById('current_rule_id');
                console.log('Found container:', container);
                
                let ruleIdInput = null;
                
                if (container) {
                    ruleIdInput = container.querySelector('input, textarea');
                    console.log('Found input element:', ruleIdInput);
                }
                
                if (ruleIdInput) {
                    console.log(`Setting rule ID input to: ${ruleId}`);
                    ruleIdInput.value = ruleId;
                    
                    // Force Gradio to recognize the change
                    ruleIdInput.dispatchEvent(new Event('input', { bubbles: true }));
                    ruleIdInput.dispatchEvent(new Event('change', { bubbles: true }));
                    
                    // Add visual feedback
                    if (action === 'preview') {
                        window.showToast('👁️ Loading preview...');
                    }
                    
                    // Trigger the appropriate button after a delay
                    setTimeout(() => {
                        let targetButton;
                        if (action === 'preview') {
                            targetButton = document.getElementById('hidden_preview_btn');
                        } else if (action === 'accept') {
                            targetButton = document.getElementById('hidden_accept_btn');
                        } else if (action === 'reject') {
                            targetButton = document.getElementById('hidden_reject_btn');
                        } else if (action === 'run') {
                            targetButton = document.getElementById('hidden_run_rule_btn');
                        } else if (action === 'archive') {
                            targetButton = document.getElementById('hidden_archive_rule_btn');
                        } else if (action === 'reactivate') {
                            targetButton = document.getElementById('hidden_reactivate_rule_btn');
                        }
                        
                        console.log(`Looking for button: hidden_${action}_btn`);
                        console.log('Found button:', targetButton);
                        
                        if (targetButton) {
                            console.log(`Clicking ${action} button`);
                            targetButton.click();
                        } else {
                            console.error(`Target button not found for action: ${action}`);
                            // List all buttons with IDs containing 'hidden' for debugging
                            const allHiddenButtons = document.querySelectorAll('[id*="hidden"]');
                            console.log('Available hidden elements:', Array.from(allHiddenButtons).map(el => el.id));
                            window.showToast(`❌ Button not found: ${action}`);
                        }
                    }, 150);
                } else {
                    console.error('Rule ID input field not found');
                    window.showToast('❌ Rule ID input not found');
                }
            } catch (error) {
                console.error('Error in handleRuleAction:', error);
                window.showToast('❌ Error: ' + error.message);
            }
        }
        
        // No need for MutationObserver - event delegation handles dynamic content
        
        // Rule card expansion
        window.toggleRuleDetails = function(ruleId) {
            const detailsDiv = document.getElementById('details-' + ruleId);
            const arrow = document.getElementById('arrow-' + ruleId);
            
            if (detailsDiv && arrow) {
                if (detailsDiv.style.display === 'none') {
                    // Collapse all other rules first
                    document.querySelectorAll('.rule-details').forEach(d => {
                        if (d.id !== 'details-' + ruleId) {
                            d.style.display = 'none';
                        }
                    });
                    document.querySelectorAll('.rule-expand-arrow').forEach(a => {
                        if (a.id !== 'arrow-' + ruleId) {
                            a.classList.remove('expanded');
                        }
                    });
                    
                    // Expand this rule
                    detailsDiv.style.display = 'block';
                    arrow.classList.add('expanded');
                } else {
                    // Collapse this rule
                    detailsDiv.style.display = 'none';
                    arrow.classList.remove('expanded');
                }
            }
        };
        
        // Demo Checklist Functions
        window.toggleChecklist = function(e) {
            if (e) {
                e.stopPropagation();
            }
            const checklist = document.getElementById('demo-checklist');
            const minimizeBtn = document.getElementById('minimize-btn');
            
            if (checklist) {
                checklist.classList.toggle('minimized');
                minimizeBtn.textContent = checklist.classList.contains('minimized') ? '▲' : '▼';
            }
        };
        
        window.updateProgress = function() {
            const checkboxes = document.querySelectorAll('.checklist-checkbox');
            const items = document.querySelectorAll('.checklist-item');
            const progressFill = document.getElementById('progress-fill');
            const progressText = document.getElementById('progress-text');
            const headerProgress = document.getElementById('header-progress');
            
            let checked = 0;
            checkboxes.forEach((cb, index) => {
                if (cb.checked) {
                    checked++;
                    items[index].classList.add('completed');
                } else {
                    items[index].classList.remove('completed');
                }
            });
            
            const percentage = (checked / checkboxes.length) * 100;
            progressFill.style.width = percentage + '%';
            progressText.textContent = `${checked}/${checkboxes.length} completed`;
            
            // Update header progress for collapsed state
            headerProgress.textContent = `(${checked}/5)`;
            
            // Check if transitioning to/from complete state
            const wasComplete = document.body.dataset.checklistComplete === 'true';
            const isNowComplete = checked === checkboxes.length;
            
            if (isNowComplete && !wasComplete) {
                // Show celebration
                showCelebration();
                document.body.dataset.checklistComplete = 'true';
            } else if (!isNowComplete && wasComplete) {
                // Hide celebration if unchecking
                dismissCelebration();
                document.body.dataset.checklistComplete = 'false';
            }
        };
        
        // Auto-show checklist in demo mode
        window.showDemoChecklist = function() {
            const checklist = document.getElementById('demo-checklist');
            if (checklist) {
                checklist.classList.add('show', 'initial-show');
                
                // Remove initial-show class after animation completes
                setTimeout(() => {
                    checklist.classList.remove('initial-show');
                }, 500);
                
                // Initialize progress display
                updateProgress();
            }
        };
        
        window.showCelebration = function() {
            const overlay = document.getElementById('celebration-overlay');
            if (overlay) {
                overlay.classList.remove('hidden');
                // Also pulse the checklist container using CSS class
                const checklist = document.getElementById('demo-checklist');
                checklist.classList.add('pulse-animation');
                setTimeout(() => {
                    checklist.classList.remove('pulse-animation');
                }, 1500);
            }
        };
        
        window.dismissCelebration = function() {
            const overlay = document.getElementById('celebration-overlay');
            if (overlay) {
                overlay.classList.add('hidden');
            }
        };
        
        // Click handler for checklist items (optional: track what users try)
        document.addEventListener('click', function(e) {
            if (e.target.classList.contains('checklist-text')) {
                const checkbox = e.target.previousElementSibling;
                if (checkbox && checkbox.type === 'checkbox') {
                    checkbox.checked = !checkbox.checked;
                    updateProgress();
                }
            }
        });
        
        // Draggable functionality for Demo Checklist
        let isDragging = false;
        let currentX;
        let currentY;
        let initialX;
        let initialY;
        let xOffset = 0;
        let yOffset = 0;
        
        window.startDrag = function(e) {
            const checklist = document.getElementById('demo-checklist');
            
            // Don't start drag if clicking the minimize button
            if (e.target.classList.contains('minimize-btn')) {
                return;
            }
            
            if (e.type === "touchstart") {
                initialX = e.touches[0].clientX - xOffset;
                initialY = e.touches[0].clientY - yOffset;
            } else {
                initialX = e.clientX - xOffset;
                initialY = e.clientY - yOffset;
            }
            
            if (e.target.classList.contains('checklist-header') || e.target.parentElement.classList.contains('checklist-header')) {
                isDragging = true;
                checklist.style.transition = 'none';
            }
        };
        
        window.dragChecklist = function(e) {
            if (isDragging) {
                e.preventDefault();
                
                if (e.type === "touchmove") {
                    currentX = e.touches[0].clientX - initialX;
                    currentY = e.touches[0].clientY - initialY;
                } else {
                    currentX = e.clientX - initialX;
                    currentY = e.clientY - initialY;
                }
                
                xOffset = currentX;
                yOffset = currentY;
                
                const checklist = document.getElementById('demo-checklist');
                checklist.style.transform = `translate(${currentX}px, ${currentY}px)`;
            }
        };
        
        window.endDrag = function(e) {
            initialX = currentX;
            initialY = currentY;
            isDragging = false;
            
            const checklist = document.getElementById('demo-checklist');
            if (checklist) {
                checklist.style.transition = '';
            }
        };
        
        // Add drag event listeners
        document.addEventListener('mousemove', dragChecklist);
        document.addEventListener('mouseup', endDrag);
        document.addEventListener('touchmove', dragChecklist);
        document.addEventListener('touchend', endDrag);
    }
    """