import streamlit as st
import base64
from io import BytesIO
import json
import time
import os
from utils import create_slide_preview, TEMPLATES, generate_slide_content
from visual_elements import search_stock_images, download_image, analyze_slide_for_visuals
from slide_editor_enhancements import render_enhanced_slide_editor

# Check for API keys
api_key = os.getenv("ANTHROPIC_API_KEY")
openai_key = os.getenv("OPENAI_API_KEY")
deepseek_key = os.getenv("DEEPSEEK_API_KEY")
perplexity_key = os.getenv("PERPLEXITY_API_KEY")
pexels_key = os.getenv("PEXELS_API_KEY")

def nav_button(label, stage, icon=None, primary=False):
    """Create a navigation button with improved styling"""
    button_style = "primary" if primary else "secondary"
    icon_text = f"{icon} " if icon else ""
    
    if st.button(f"{icon_text}{label}", key=f"nav_{stage}", type=button_style, use_container_width=True):
        st.session_state.current_stage = stage
        st.rerun()

def display_navigation_bar():
    """Display navigation bar at the top of the page"""
    stages = [
        {"name": "ideation", "label": "Ideation", "icon": "💡"},
        {"name": "storyboard", "label": "Storyboard", "icon": "📋"},
        {"name": "template", "label": "Template", "icon": "🎨"},
        {"name": "slides", "label": "Edit Slides", "icon": "🖼️"},
        {"name": "export", "label": "Export", "icon": "📤"}
    ]
    
    cols = st.columns(len(stages))
    
    for i, stage in enumerate(stages):
        with cols[i]:
            is_current = st.session_state.current_stage == stage["name"]
            
            # Create a clickable button with custom CSS
            if st.button(
                f"{stage['icon']} {stage['label']}", 
                key=f"nav_top_{stage['name']}", 
                disabled=is_current,
                use_container_width=True,
                type="primary" if is_current else "secondary"
            ):
                # If we've already completed previous stages, allow direct navigation
                if stage["name"] in ["ideation"] or (
                    stage["name"] == "storyboard" and st.session_state.storyboard
                ) or (
                    stage["name"] == "template" and st.session_state.storyboard
                ) or (
                    stage["name"] == "slides" and st.session_state.slides_content
                ) or (
                    stage["name"] == "export" and st.session_state.slides_content
                ):
                    st.session_state.current_stage = stage["name"]
                    st.rerun()
                else:
                    st.warning(f"Please complete the previous stages before going to {stage['label']}")
    
    # Progress bar
    stages_order = ["ideation", "storyboard", "template", "slides", "export"]
    current_idx = stages_order.index(st.session_state.current_stage)
    progress = current_idx / (len(stages_order) - 1)
    st.progress(progress)

def render_ai_settings():
    """Render AI settings sidebar section"""
    st.sidebar.write("## 🧠 AI Settings")
    
    # Get AI manager
    try:
        from multi_llm_provider import get_ai_manager
        ai_manager = get_ai_manager()
        available_models = ai_manager.get_available_models()
    except ImportError:
        available_models = {
            "claude-3-sonnet-20250219": "Claude 3 Sonnet",
            "claude-3-haiku-20250319": "Claude 3 Haiku",
            "claude-3-opus-20250229": "Claude 3 Opus"
        }
    
    if not available_models:
        st.sidebar.warning("No AI providers configured. Add API keys in secrets.")
        default_model = "claude-3-sonnet-20250219"
    else:
        default_model = list(available_models.keys())[0] if available_models else "claude-3-sonnet-20250219"
    
    # Select default model for the whole presentation
    selected_model = st.sidebar.selectbox(
        "Default AI Model",
        options=list(available_models.keys()),
        format_func=lambda x: available_models.get(x, x),
        index=0 if default_model in available_models else 0
    )
    
    st.session_state.default_model = selected_model
    
    # Temperature setting
    temperature = st.sidebar.slider(
        "AI Creativity",
        min_value=0.0,
        max_value=1.0,
        value=0.7,
        step=0.1,
        help="Higher values make output more creative but less predictable"
    )
    
    st.session_state.ai_temperature = temperature
    
    # Web search integration
    enable_search = st.sidebar.checkbox(
        "Enable Web Search",
        value=st.session_state.get("enable_web_search", False),
        help="Use Perplexity to search for up-to-date information",
        disabled=not perplexity_key
    )
    
    st.session_state.enable_web_search = enable_search
    
    if enable_search and not perplexity_key:
        st.sidebar.warning("Perplexity API key required for web search")
    
    # Stock image settings
    st.sidebar.write("## 🖼️ Image Settings")
    
    if not pexels_key:
        st.sidebar.warning("Pexels API key required for better stock images")

def render_ideation_stage():
    """Render the ideation stage UI"""
    display_navigation_bar()
    
    # Add AI settings sidebar
    render_ai_settings()
    
    st.header("💡 Step 1: Presentation Ideation")
    
    st.write("Let's start by defining the purpose and scope of your presentation.")
    
    col1, col2 = st.columns(2)
    
    with col1:
        st.session_state.presentation_title = st.text_input(
            "Presentation Title", 
            value=st.session_state.presentation_title
        )
        
        st.session_state.presentation_purpose = st.text_area(
            "What's the purpose of this presentation?",
            value=st.session_state.presentation_purpose,
            help="E.g., Inform, persuade, pitch a product, update stakeholders, etc."
        )
    
    with col2:
        st.session_state.target_audience = st.text_area(
            "Who is your target audience?",
            value=st.session_state.target_audience,
            help="E.g., Executives, customers, technical team, general public, etc."
        )
        
        # Add example suggestions
        st.write("📚 **Example presentations:**")
        examples = [
            "Quarterly Business Review",
            "Product Launch Presentation",
            "Team Project Update",
            "Investor Pitch Deck"
        ]
        
        # Create two columns for examples
        ex_col1, ex_col2 = st.columns(2)
        
        for i, example in enumerate(examples):
            with ex_col1 if i % 2 == 0 else ex_col2:
                if st.button(f"📄 {example}", key=f"example_{example}"):
                    st.session_state.presentation_title = example
                    
                    # Set appropriate purpose and audience based on example
                    if "Quarterly" in example:
                        st.session_state.presentation_purpose = "Review business performance for the past quarter and outline goals for the next quarter"
                        st.session_state.target_audience = "Executives and department heads"
                    elif "Product" in example:
                        st.session_state.presentation_purpose = "Introduce a new product to customers and highlight its key features and benefits"
                        st.session_state.target_audience = "Potential customers and sales team"
                    elif "Project" in example:
                        st.session_state.presentation_purpose = "Update team members on project progress, achievements, and next steps"
                        st.session_state.target_audience = "Project stakeholders and team members"
                    elif "Investor" in example:
                        st.session_state.presentation_purpose = "Pitch our business to potential investors to secure funding"
                        st.session_state.target_audience = "Venture capitalists and angel investors"
                    st.rerun()
    
    # AI model selection for storyboard generation
    st.write("### 🧠 AI Engine Selection")
    
    # Get available models
    try:
        from multi_llm_provider import get_ai_manager
        ai_manager = get_ai_manager()
        available_models = ai_manager.get_available_models()
    except (ImportError, Exception):
        available_models = {
            "claude-3-sonnet-20250219": "Claude 3 Sonnet",
            "claude-3-haiku-20250319": "Claude 3 Haiku",
            "claude-3-opus-20250229": "Claude 3 Opus"
        }
    
    cols = st.columns([2, 1])
    with cols[0]:
        selected_model = st.selectbox(
            "AI Model for Storyboard Generation",
            options=list(available_models.keys()),
            format_func=lambda x: available_models.get(x, x),
            index=0,
            key="storyboard_model"
        )
    
    with cols[1]:
        web_research = st.checkbox(
            "Include Web Research",
            value=st.session_state.get("enable_web_search", False),
            disabled=not perplexity_key,
            help="Search the web for the latest information related to your presentation topic"
        )
    
    # Generate button with loading indicator
    st.markdown("---")
    if st.button("🐊 Generate Storyboard", type="primary", use_container_width=True):
        if not st.session_state.presentation_title or not st.session_state.presentation_purpose:
            st.warning("Please provide at least a title and purpose to proceed.")
        else:
            with st.spinner("🧠 SlideGator.AI is chomping on your presentation ideas..."):
                from utils import generate_storyboard
                
                # Generate storyboard 
                storyboard = generate_storyboard(
                    st.session_state.presentation_title,
                    st.session_state.presentation_purpose,
                    st.session_state.target_audience,
                    model=selected_model
                )
                
                if storyboard:
                    st.session_state.storyboard = storyboard
                    st.session_state.current_stage = "storyboard"
                    st.success("Storyboard created successfully!")
                    st.rerun()

def render_storyboard_stage():
    """Render the storyboard review stage UI (simplified version without AI features)"""
    display_navigation_bar()
    st.header("📋 Step 2: Review and Edit Storyboard")
    
    st.write(f"Storyboard for: **{st.session_state.presentation_title}**")
    st.write("Expand each slide to edit its content or add notes.")
    
    # Display storyboard for review with edit options
    edited_storyboard = []
    
    # Simple regenerate button 
    if st.button("🔄 Regenerate All", help="Create a new storyboard with current title/purpose"):
        st.info("Storyboard regeneration feature coming soon")
    
    # Add slide button with options
    col1, col2 = st.columns([3, 1])
    with col1:
        slide_position = st.radio(
            "Add new slide:",
            ["At End", "After Current", "Before Current"],
            horizontal=True
        )
    
    with col2:
        if st.button("➕ Add New Slide", use_container_width=True):
            # Create a new slide
            new_slide = {
                'title': 'New Slide',
                'purpose': 'Additional content',
                'key_points': ['Add your content here'],
                'visual_elements': ['Suggested visuals will appear here']
            }
            
            # Insert at the selected position
            if slide_position == "At End":
                st.session_state.storyboard.append(new_slide)
            elif slide_position == "After Current":
                # Get current slide index from session state or default to last
                current_idx = st.session_state.get("current_storyboard_slide", len(st.session_state.storyboard) - 1)
                st.session_state.storyboard.insert(current_idx + 1, new_slide)
            else:  # Before Current
                current_idx = st.session_state.get("current_storyboard_slide", 0)
                st.session_state.storyboard.insert(max(0, current_idx), new_slide)
            
            st.rerun()
    
    # Display storyboard slides
    for i, slide in enumerate(st.session_state.storyboard):
        # Set current slide in session state when expander is opened
        is_expanded = i == 0 or st.session_state.get("current_storyboard_slide") == i
        
        with st.expander(f"Slide {i+1}: {slide.get('title', 'Untitled')}", expanded=is_expanded):
            if is_expanded:
                st.session_state.current_storyboard_slide = i
            
            # Main slide content
            cols = st.columns([3, 1])
            with cols[0]:
                slide_title = st.text_input(f"Title ###{i}", value=slide.get('title', 'Untitled'))
            with cols[1]:
                # Slide reordering and deletion
                cols2 = st.columns([1, 1, 1])
                with cols2[0]:
                    if i > 0 and st.button("⬆️", key=f"up_{i}"):
                        st.session_state.storyboard[i], st.session_state.storyboard[i-1] = st.session_state.storyboard[i-1], st.session_state.storyboard[i]
                        st.session_state.current_storyboard_slide = i - 1
                        st.rerun()
                with cols2[1]:
                    if i < len(st.session_state.storyboard) - 1 and st.button("⬇️", key=f"down_{i}"):
                        st.session_state.storyboard[i], st.session_state.storyboard[i+1] = st.session_state.storyboard[i+1], st.session_state.storyboard[i]
                        st.session_state.current_storyboard_slide = i + 1
                        st.rerun()
                with cols2[2]:
                    if st.button("🗑️", key=f"delete_{i}"):
                        if len(st.session_state.storyboard) > 1:  # Prevent deleting the last slide
                            st.session_state.storyboard.pop(i)
                            st.session_state.current_storyboard_slide = min(i, len(st.session_state.storyboard) - 1)
                            st.rerun()
                        else:
                            st.error("Cannot delete the last slide")
            
            slide_purpose = st.text_area(f"Purpose ###{i}", value=slide.get('purpose', ''))
            
            # Handle key points (could be string or list)
            if isinstance(slide.get('key_points', ""), list):
                key_points_text = "\n".join(slide['key_points'])
            else:
                key_points_text = slide.get('key_points', "")
                
            key_points = st.text_area(f"Key Points (one per line) ###{i}", value=key_points_text)
            
            # Handle visual elements (could be string or list)
            if isinstance(slide.get('visual_elements', ""), list):
                visual_elements_text = "\n".join(slide['visual_elements'])
            else:
                visual_elements_text = slide.get('visual_elements', "")
                
            visual_elements = st.text_area(f"Visual Elements (one per line) ###{i}", value=visual_elements_text)
            
            # Update storyboard with edits
            edited_slide = {
                'title': slide_title,
                'purpose': slide_purpose,
                'key_points': key_points.split("\n") if "\n" in key_points else [key_points] if key_points else [],
                'visual_elements': visual_elements.split("\n") if "\n" in visual_elements else [visual_elements] if visual_elements else []
            }
            edited_storyboard.append(edited_slide)
    
    # Update the storyboard in session state
    st.session_state.storyboard = edited_storyboard
    
    st.markdown("---")
    col1, col2 = st.columns(2)
    with col1:
        nav_button("Back to Ideation", "ideation", icon="⬅️")
    with col2:
        nav_button("Continue to Template Selection", "template", icon="➡️", primary=True)

def render_template_stage():
    """Render the template selection stage UI with enhanced options"""
    display_navigation_bar()
    
    # Add AI settings sidebar
    render_ai_settings()
    
    st.header("🎨 Step 3: Select a Template")
    
    st.write("Choose a visual style for your presentation:")
    
    # Preview section
    preview_col, options_col = st.columns([2, 1])
    
    with options_col:
        # Add template upload option
        st.subheader("Custom Template")
        st.write("Upload your own PowerPoint template:")
        uploaded_template = st.file_uploader("Upload PPTX template", type=["pptx"])
        
        if uploaded_template:
            # Save the uploaded template
            from utils import add_custom_template
            success, message = add_custom_template(uploaded_template)
            if success:
                st.success(message)
                st.session_state.selected_template = "custom"
                # Store template file name for display
                st.session_state.custom_template_name = uploaded_template.name
            else:
                st.error(message)
        
        # Template color customization
        st.subheader("Color Customization")
        
        # Get current template
        current_template = st.session_state.selected_template
        template_info = TEMPLATES.get(current_template, TEMPLATES["professional"])
        colors = template_info["colors"]
        
        # Allow color customization
        primary_color = st.color_picker(
            "Primary Color", 
            value=colors.get("primary", "#0F52BA"),
            key="primary_color"
        )
        
        accent_color = st.color_picker(
            "Accent Color", 
            value=colors.get("accent", "#D4AF37"),
            key="accent_color"
        )
        
        # Store customized colors
        if "custom_colors" not in st.session_state:
            st.session_state.custom_colors = {}
        
        st.session_state.custom_colors[current_template] = {
            "primary": primary_color,
            "accent": accent_color
        }
        
        # Apply customized colors
        if st.button("Apply Custom Colors"):
            # Update the template colors
            if current_template in TEMPLATES:
                TEMPLATES[current_template]["colors"]["primary"] = primary_color
                TEMPLATES[current_template]["colors"]["accent"] = accent_color
                st.success("Custom colors applied!")
    
    with preview_col:
        # Template preview - show a sample slide with the selected template
        st.subheader("Template Preview")
        
        # Create a sample slide for preview
        sample_slide = {
            "title": "Sample Slide",
            "content": [
                "This is how your slides will look",
                "With the selected template style",
                "You can customize colors and fonts"
            ],
            "visual_elements": ["Chart showing data trends", "Icon representing growth"]
        }
        
        # Get current template
        current_template = st.session_state.selected_template
        
        # Show preview based on template
        try:
            from visual_elements import generate_html_preview_with_visuals
            preview_html = generate_html_preview_with_visuals(sample_slide, current_template)
            st.components.v1.html(preview_html, height=300)
        except Exception as e:
            st.error(f"Error generating preview: {str(e)}")
            # Fall back to standard preview
            preview_html = create_slide_preview(sample_slide, current_template)
            st.components.v1.html(preview_html, height=300)
        
        # Display current template name
        if current_template == "custom" and "custom_template_name" in st.session_state:
            st.info(f"Using custom template: {st.session_state.custom_template_name}")
        else:
            st.info(f"Currently using: {current_template.title()} template")
    
    st.markdown("---")
    
    # Create template cards in a grid
    st.subheader("🐊 SlideGator Template Gallery")
    st.write("Select from our professionally designed templates:")
    
    # Use columns to create a grid
    cols = st.columns(3)
    
    # Add Professional template
    with cols[0]:
        st.subheader("Professional")
        st.write("Clean, corporate style with blue and gray.")
        st.markdown("""
        <div style="border:1px solid #ddd; padding:10px; border-radius:5px; background-color:#f8f9fa;">
            <div style="height:20px; background-color:#0F52BA;"></div>
            <div style="padding:10px; text-align:center;">
                <div style="font-weight:bold; margin-bottom:10px;">Professional Style</div>
                <div style="height:5px; background-color:#D4AF37; width:50%; margin:auto;"></div>
                <div style="font-size:0.8em; margin-top:10px;">Clear and business-focused</div>
            </div>
        </div>
        """, unsafe_allow_html=True)
        if st.button("🎯 Select Professional", key="select_prof"):
            st.session_state.selected_template = "professional"
            st.success("Selected: Professional template")
            st.rerun()
    
    # Add Creative template
    with cols[1]:
        st.subheader("Creative")
        st.write("Vibrant design with modern elements.")
        st.markdown("""
        <div style="border:1px solid #ddd; padding:10px; border-radius:5px; background-color:#FFDE59;">
            <div style="height:20px; background-color:#FF5757;"></div>
            <div style="padding:10px; text-align:center;">
                <div style="font-weight:bold; margin-bottom:10px;">Creative Style</div>
                <div style="height:5px; background-color:#5CE1E6; width:50%; margin:auto;"></div>
                <div style="font-size:0.8em; margin-top:10px;">Bold and engaging design</div>
            </div>
        </div>
        """, unsafe_allow_html=True)
        if st.button("🎨 Select Creative", key="select_creative"):
            st.session_state.selected_template = "creative"
            st.success("Selected: Creative template")
            st.rerun()
    
    # Add Minimalist template
    with cols[2]:
        st.subheader("Minimalist")
        st.write("Simple design with lots of whitespace.")
        st.markdown("""
        <div style="border:1px solid #ddd; padding:10px; border-radius:5px; background-color:#FFFFFF;">
            <div style="height:20px; background-color:#000000;"></div>
            <div style="padding:10px; text-align:center;">
                <div style="font-weight:bold; margin-bottom:10px;">Minimalist Style</div>
                <div style="height:1px; background-color:#000000; width:50%; margin:auto;"></div>
                <div style="font-size:0.8em; margin-top:10px;">Clean and elegant</div>
            </div>
        </div>
        """, unsafe_allow_html=True)
        if st.button("✨ Select Minimalist", key="select_min"):
            st.session_state.selected_template = "minimalist"
            st.success("Selected: Minimalist template")
            st.rerun()
    
    # Font selection
    with st.expander("🔤 Font Settings"):
        st.write("Choose font styles for your presentation:")
        
        font_options = [
            "Arial, sans-serif", 
            "Helvetica, sans-serif", 
            "Calibri, sans-serif", 
            "Georgia, serif", 
            "Times New Roman, serif", 
            "Verdana, sans-serif", 
            "Tahoma, sans-serif"
        ]
        
        font_cols = st.columns(2)
        
        with font_cols[0]:
            title_font = st.selectbox(
                "Title Font",
                options=font_options,
                index=font_options.index(TEMPLATES[current_template]["fonts"]["title"]) if TEMPLATES[current_template]["fonts"]["title"] in font_options else 0
            )
        
        with font_cols[1]:
            body_font = st.selectbox(
                "Body Font",
                options=font_options,
                index=font_options.index(TEMPLATES[current_template]["fonts"]["body"]) if TEMPLATES[current_template]["fonts"]["body"] in font_options else 0
            )
        
        if st.button("Apply Font Settings"):
            # Update the template fonts
            if current_template in TEMPLATES:
                TEMPLATES[current_template]["fonts"]["title"] = title_font
                TEMPLATES[current_template]["fonts"]["body"] = body_font
                st.success("Font settings applied!")
    
    st.markdown("---")
    col1, col2 = st.columns(2)
    with col1:
        nav_button("Back to Storyboard", "storyboard", icon="⬅️")
    with col2:
        if st.button("🐊 Generate Slides ➡️", type="primary", use_container_width=True):
            st.session_state.current_stage = "slides"
            # Generate detailed content for each slide
            with st.spinner("🐊 SlideGator.AI is crafting your slides..."):
                slides_content = []
                progress_placeholder = st.empty()
                
                for i, slide in enumerate(st.session_state.storyboard):
                    # Calculate progress as a value between 0 and 1
                    progress = i / len(st.session_state.storyboard)
                    progress_placeholder.progress(progress)
                    progress_placeholder.text(f"Generating slide {i+1} of {len(st.session_state.storyboard)}: {slide.get('title', 'Untitled')}")
                    
                    # Get model for slide generation
                    model = slide.get("ai_settings", {}).get("model", st.session_state.get("default_model", "claude-3-sonnet-20250219"))
                    
                    slide_content = generate_slide_content(slide, st.session_state.selected_template, model=model)
                    slides_content.append(slide_content)
                
                progress_placeholder.progress(1.0)
                progress_placeholder.empty()
                st.session_state.slides_content = slides_content
                st.success("All slides generated!")
                st.rerun()

def render_slides_stage():
    """Render the enhanced slide editing stage UI"""
    display_navigation_bar()
    
    # Add AI settings sidebar
    render_ai_settings()
    
    st.header("🖼️ Step 4: Review and Edit Slides")
    
    if not st.session_state.slides_content:
        st.warning("No slides content available. Please go back and generate slides.")
        nav_button("Back to Template Selection", "template", icon="⬅️")
        return
    
    # Add slide navigation controls
    col1, col2, col3 = st.columns([1, 2, 1])
    with col1:
        prev_slide = st.button("◀️ Previous Slide")
    with col2:
        current_slide = st.session_state.get("current_slide_index", 0)
        slide_selector = st.select_slider(
            "Navigate Slides",
            options=list(range(len(st.session_state.slides_content))),
            value=current_slide,
            format_func=lambda x: f"Slide {x+1}: {st.session_state.slides_content[x].get('title', 'Untitled')}"
        )
        st.session_state.current_slide_index = slide_selector
    with col3:
        next_slide = st.button("Next Slide ▶️")
    
    if prev_slide and st.session_state.current_slide_index > 0:
        st.session_state.current_slide_index -= 1
        st.rerun()
    
    if next_slide and st.session_state.current_slide_index < len(st.session_state.slides_content) - 1:
        st.session_state.current_slide_index += 1
        st.rerun()
    
    # Show slide editing interface for the current slide
    current_slide_idx = st.session_state.current_slide_index
    current_slide = st.session_state.slides_content[current_slide_idx]
    
    # Add slide management controls
    col1, col2, col3, col4 = st.columns([1, 1, 1, 1])
    with col1:
        if current_slide_idx > 0 and st.button("⬆️ Move Up"):
            st.session_state.slides_content[current_slide_idx], st.session_state.slides_content[current_slide_idx-1] = st.session_state.slides_content[current_slide_idx-1], st.session_state.slides_content[current_slide_idx]
            st.session_state.current_slide_index -= 1
            st.rerun()
    
    with col2:
        if current_slide_idx < len(st.session_state.slides_content) - 1 and st.button("⬇️ Move Down"):
            st.session_state.slides_content[current_slide_idx], st.session_state.slides_content[current_slide_idx+1] = st.session_state.slides_content[current_slide_idx+1], st.session_state.slides_content[current_slide_idx]
            st.session_state.current_slide_index += 1
            st.rerun()
    
    with col3:
        if st.button("➕ Add Slide"):
            # Create a new slide based on the current one
            new_slide = {
                "title": "New Slide",
                "content": ["Add your content here"],
                "visual_elements": ["Add visual elements here"],
                "notes": "Add presenter notes here"
            }
            # Insert after current slide
            st.session_state.slides_content.insert(current_slide_idx + 1, new_slide)
            st.session_state.current_slide_index += 1
            st.rerun()
    
    with col4:
        if len(st.session_state.slides_content) > 1 and st.button("🗑️ Delete Slide"):
            st.session_state.slides_content.pop(current_slide_idx)
            if current_slide_idx >= len(st.session_state.slides_content):
                st.session_state.current_slide_index = len(st.session_state.slides_content) - 1
            st.rerun()
    
    # Current template
    template_name = st.session_state.selected_template
    
    # Editor for current slide - use enhanced slide editor
    st.write(f"### Editing Slide {current_slide_idx + 1}")
    updated_slide = render_enhanced_slide_editor(current_slide_idx, current_slide, template_name)
    st.session_state.slides_content[current_slide_idx] = updated_slide
    
    # Navigation buttons at bottom
    st.markdown("---")
    col1, col2 = st.columns(2)
    with col1:
        nav_button("Back to Template Selection", "template", icon="⬅️")
    with col2:
        nav_button("Finalize Presentation", "export", icon="➡️", primary=True)

def render_export_stage():
    """Render the enhanced export stage UI"""
    display_navigation_bar()
    
    # Add AI settings sidebar
    render_ai_settings()
    
    st.header("📤 Step 5: Export Presentation")
    
    st.write("Your presentation is ready to export!")
    
    # Add final quality check option
    with st.expander("🔍 Quality Check", expanded=True):
        st.write("Run a final quality check on your presentation before exporting:")
        
        quality_options = st.multiselect(
            "Select checks to run:",
            ["Spelling & Grammar", "Visual Balance", "Content Consistency", "Presentation Flow"],
            default=["Spelling & Grammar", "Content Consistency"]
        )
        
        if st.button("Run Quality Check", use_container_width=True):
            with st.spinner("🔍 Running quality checks..."):
                # Show progress
                progress = st.progress(0)
                
                for i, check in enumerate(quality_options):
                    # Update progress
                    progress.progress((i + 0.5) / len(quality_options))
                    
                    if check == "Spelling & Grammar":
                        with st.spinner("Checking spelling and grammar..."):
                            # In a real implementation, this would check all slide content
                            time.sleep(1)  # Simulate processing
                            st.success("✅ Spelling and grammar check complete. No major issues found.")
                    
                    elif check == "Visual Balance":
                        with st.spinner("Analyzing visual balance..."):
                            # Simulate visual analysis
                            time.sleep(1)
                            st.info("ℹ️ Visual recommendation: Consider adding more images to slides 3 and 5.")
                    
                    elif check == "Content Consistency":
                        with st.spinner("Checking content consistency..."):
                            # In a real implementation, would use AI to check consistency
                            time.sleep(1)
                            st.success("✅ Content is consistent throughout the presentation.")
                    
                    elif check == "Presentation Flow":
                        with st.spinner("Analyzing presentation flow..."):
                            # Simulate flow analysis
                            time.sleep(1)
                            st.info("ℹ️ Flow recommendation: Consider adding a transition slide between slides 2 and 3.")
                
                # Complete progress
                progress.progress(1.0)
    
    # Export options
    export_tabs = st.tabs(["PowerPoint", "PDF", "Images", "Advanced"])
    
    with export_tabs[0]:
        # PowerPoint export
        st.subheader("PowerPoint Export")
        
        # Create the PowerPoint file
        with st.spinner("🐊 SlideGator.AI is finalizing your presentation..."):
            try:
                try:
                    # Try to use enhanced PowerPoint creation if available
                    from enhanced_pptx import enhanced_create_ppt
                    ppt_buffer = enhanced_create_ppt(
                        st.session_state.slides_content,
                        st.session_state.selected_template
                    )
                except ImportError:
                    # Fall back to standard PowerPoint creation
                    from utils import create_ppt
                    ppt_buffer = create_ppt(
                        st.session_state.slides_content,
                        st.session_state.selected_template
                    )
                creation_success = True
            except Exception as e:
                st.error(f"Error creating PowerPoint file: {str(e)}")
                import traceback
                st.error(traceback.format_exc())
                creation_success = False
        
        if creation_success:
            # Provide download link
            filename = f"{st.session_state.presentation_title.replace(' ', '_')}.pptx"
            
            # Get download link HTML
            b64 = base64.b64encode(ppt_buffer.read()).decode()
            href = f'<a href="data:application/vnd.openxmlformats-officedocument.presentationml.presentation;base64,{b64}" download="{filename}" class="download-btn">🐊 Download PowerPoint Presentation</a>'
            
            # Add some styling to the download button
            st.markdown("""
            <style>
            .download-btn {
                display: inline-block;
                padding: 12px 20px;
                background-color: #4CAF50;
                color: white;
                text-decoration: none;
                border-radius: 4px;
                font-weight: bold;
                text-align: center;
                transition: background-color 0.3s;
                font-size: 1.2em;
            }
            .download-btn:hover {
                background-color: #45a049;
                box-shadow: 0 4px 8px rgba(0,0,0,0.2);
            }
            </style>
            """, unsafe_allow_html=True)
            
            # Display download link
            st.markdown(f"<div style='text-align: center; margin: 30px 0;'>{href}</div>", unsafe_allow_html=True)
    
    with export_tabs[1]:
        # PDF export
        st.subheader("PDF Export")
        st.info("PDF export feature is coming soon. For now, please use the PowerPoint export and save as PDF.")
    
    with export_tabs[2]:
        # Image export
        st.subheader("Image Export")
        st.info("Image export feature is coming soon. This will allow you to export each slide as a PNG or JPG.")
    
    with export_tabs[3]:
        # Advanced export options
        st.subheader("Advanced Export Options")
        
        # Customization before export
        st.write("Customize final export settings:")
        
        optimize_file = st.checkbox("Optimize file size", value=True)
        include_notes = st.checkbox("Include presenter notes", value=True)
        protect_file = st.checkbox("Add password protection", value=False)
        
        if protect_file:
            password = st.text_input("Set password (optional)", type="password")
        
        st.info("Advanced export features are coming soon. These settings are for demonstration purposes only.")
    
    # Display congratulations message
    st.markdown("""
    <div style="text-align:center; padding: 20px; margin: 20px 0; border-radius: 5px; background-color: #f8f9fa;">
        <div style="font-size: 2rem; margin-bottom: 10px;">🎉 Congratulations! 🎉</div>
        <div style="font-size: 1.2rem; margin-bottom: 15px;">Your presentation is ready!</div>
        <div style="font-size: 1rem; color: #666;">SlideGator.AI has snapped up the perfect presentation for you.</div>
    </div>
    """, unsafe_allow_html=True)
    
    # Preview of slides
    st.subheader("📋 Presentation Preview")
    
    # Display a sample of slides
    preview_cols = st.columns(3)
    for i, slide in enumerate(st.session_state.slides_content[:3]):  # Just show first 3 slides
        with preview_cols[i % 3]:
            try:
                # Use enhanced preview with visuals if available
                try:
                    from visual_elements import generate_html_preview_with_visuals
                    preview_html = generate_html_preview_with_visuals(slide, st.session_state.selected_template)
                except ImportError:
                    preview_html = create_slide_preview(slide, st.session_state.selected_template)
                st.components.v1.html(preview_html, height=200)
                st.write(f"**Slide {i+1}:** {slide.get('title', 'Untitled')}")
            except Exception as e:
                st.error(f"Error previewing slide {i+1}")
    
    # Show "View All" option
    if len(st.session_state.slides_content) > 3:
        st.write("...")
        show_all = st.checkbox("Show all slides", value=False)
        
        if show_all:
            st.write("### All Slides")
            for i, slide in enumerate(st.session_state.slides_content):
                with st.expander(f"Slide {i+1}: {slide.get('title', 'Untitled')}"):
                    try:
                        # Try enhanced preview
                        try:
                            from visual_elements import generate_html_preview_with_visuals
                            preview_html = generate_html_preview_with_visuals(slide, st.session_state.selected_template)
                        except ImportError:
                            preview_html = create_slide_preview(slide, st.session_state.selected_template)
                        st.components.v1.html(preview_html, height=300)
                    except Exception as e:
                        st.error(f"Error previewing slide {i+1}")
    
    # Navigation buttons
    st.markdown("---")
    col1, col2 = st.columns(2)
    with col1:
        nav_button("Back to Edit Slides", "slides", icon="⬅️")
    with col2:
        if st.button("🐊 Start New Presentation", type="primary", use_container_width=True):
            # Ask for confirmation
            if st.checkbox("✅ Confirm starting a new presentation? All current work will be lost.", value=False):
                # Reset session state
                for key in list(st.session_state.keys()):
                    if key not in ["session_id", "ai_manager"]:
                        del st.session_state[key]
                
                st.session_state.current_stage = "ideation"
                st.session_state.presentation_title = ""
                st.session_state.presentation_purpose = ""
                st.session_state.target_audience = ""
                st.session_state.storyboard = []
                st.session_state.selected_template = "professional"
                st.session_state.slides_content = []
                
                st.success("Starting new presentation...")
                st.rerun()