import gradio as gr
import os
import glob
from pathlib import Path

def create_flux_krea_gallery():
    """Create the FLUX.1-Krea-dev image gallery"""
    # Simulate loading the 100 images (in a real app, these would be actual files)
    image_files = []
    for i in range(370, 470):
        # Format the filename with leading zeros
        filename = f"flux_krea_{i:05d}_.png"
        image_files.append((f"images/{filename}", f"FLUX.1-Krea-dev Image {i}"))
    
    return image_files

def get_comfyui_workflow_image():
    """Get the ComfyUI workflow screenshot"""
    return "workflows/Screenshot 2025-10-28 at 5.00.24 PM.png"

def load_comfyui_log():
    """Load and display the ComfyUI log content"""
    try:
        with open("logs/comfyui_8000.log", "r", encoding="utf-8") as file:
            log_content = file.read()
        return log_content
    except FileNotFoundError:
        return "Log file 'logs/comfyui_8000.log' not found. Please ensure the log file exists in the logs directory."
    except Exception as e:
        return f"Error reading log file: {str(e)}"

def create_demo():
    """Create the main Gradio demo"""
    
    with gr.Blocks(
        title="FLUX.1-Krea-dev & ComfyUI MUSTARD Showcase",
        theme=gr.themes.Soft(
            primary_hue="blue",
            secondary_hue="purple",
        )
    ) as demo:
        
        # Header
        gr.Markdown(
            """
            # 🎨 FLUX.1-Krea-dev & ComfyUI MUSTARD Showcase
            *Advanced AI Image Generation Workflow Demonstration*
            """
        )
        
        # Main tabs
        with gr.Tabs():
            
            # Tab 1: FLUX.1-Krea-dev Gallery
            with gr.TabItem("🖼️ FLUX.1-Krea-dev MUSTARD Images"):
                gr.Markdown(
                    """
                    ## FLUX.1-Krea-dev Generated Images
                    *Showing all 100 generated images from flux_krea_00370_.png to flux_krea_00469_.png*
                    """
                )
                
                # Image gallery
                gallery = gr.Gallery(
                    label="FLUX.1-Krea-dev Image Collection",
                    elem_id="gallery",
                    columns=5,
                    rows=5,
                    height="auto",
                    object_fit="contain"
                )
                
                # Load gallery button
                with gr.Row():
                    load_gallery_btn = gr.Button("🔄 Load FLUX.1-Krea-dev Gallery", variant="primary")
                    clear_gallery_btn = gr.Button("🗑️ Clear Gallery", variant="secondary")
                
                # Gallery controls
                with gr.Row():
                    gr.Markdown("**Gallery Controls:** Use the arrows to navigate, click on images to view them in full size.")
            
            # Tab 2: ComfyUI Workflow
            with gr.TabItem("⚙️ ComfyUI Workflow"):
                gr.Markdown(
                    """
                    ## ComfyUI Workflow Configuration
                    *Workflow screenshot: `Screenshot 2025-10-28 at 5.00.24 PM.png`*
                    """
                )
                
                # Workflow image display
                workflow_image = gr.Image(
                    label="ComfyUI Workflow",
                    elem_id="workflow-image",
                    height=600,
                    interactive=False
                )
                
                # Workflow description
                gr.Markdown(
                    """
                    ### Workflow Description
                    This ComfyUI workflow demonstrates the FLUX.1-Krea-dev model configuration with:
                    - **Load Checkpoint**: FLUX.1-Krea-dev model
                    - **CLIP Text Encode**: Positive and negative prompts
                    - **KSampler**: Euler, 20 steps
                    - **VAE Decode**: Convert latent to image
                    - **Save Image**: Output to specified directory
                    """
                )
            
            # Tab 3: ComfyUI Log
            with gr.TabItem("📋 ComfyUI Log"):
                gr.Markdown(
                    """
                    ## ComfyUI Generation Log
                    *Log file: `comfyui_8000.log` - Complete image generation process*
                    """
                )
                
                # Log display
                log_display = gr.Textbox(
                    label="ComfyUI Log Content",
                    elem_id="log-display",
                    lines=25,
                    max_lines=50,
                    show_copy_button=True,
                    autoscroll=False
                )
                
                # Log controls
                with gr.Row():
                    load_log_btn = gr.Button("📄 Load ComfyUI Log", variant="primary")
                    clear_log_btn = gr.Button("🗑️ Clear Log", variant="secondary")
        
        # Function bindings
        load_gallery_btn.click(
            fn=create_flux_krea_gallery,
            outputs=gallery
        )
        
        clear_gallery_btn.click(
            fn=lambda: [],
            outputs=gallery
        )
        
        load_log_btn.click(
            fn=load_comfyui_log,
            outputs=log_display
        )
        
        clear_log_btn.click(
            fn=lambda: "",
            outputs=log_display
        )
        
        # Initialize with sample data
        demo.load(
            fn=create_flux_krea_gallery,
            outputs=gallery
        )
        
        demo.load(
            fn=load_comfyui_log,
            outputs=log_display
        )
        
        # Add some CSS for better styling
        demo.css = """
        #gallery {
            min-height: 600px;
        }
        #workflow-image {
            border: 2px solid #e2e8f0;
            border-radius: 10px;
        }
        #log-display {
            font-family: 'Monaco', 'Menlo', 'Ubuntu Mono', monospace;
            font-size: 12px;
        }
        .tab-nav {
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        }
        """
    
    return demo

def create_advanced_demo():
    """Create an enhanced version with more features"""
    
    with gr.Blocks(
        title="FLUX.1-Krea-dev Advanced Showcase",
        theme=gr.themes.Glass(
            primary_hue="purple",
            secondary_hue="blue",
        )
    ) as demo:
        
        # Custom CSS
        demo.css = """
        .header-banner {
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            padding: 2rem;
            border-radius: 15px;
            color: white;
            margin-bottom: 2rem;
        }
        .stat-card {
            background: white;
            padding: 1rem;
            border-radius: 10px;
            box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
            margin: 0.5rem;
        }
        """
        
        # Header Banner
        gr.HTML("""
        <div class="header-banner">
            <h1 style="margin: 0; font-size: 2.5em;">🎨 FLUX.1-Krea-dev & ComfyUI</h1>
            <p style="margin: 0; font-size: 1.2em; opacity: 0.9;">Advanced AI Image Generation Workflow Showcase</p>
        </div>
        """)
        
        # Statistics Row
        with gr.Row():
            with gr.Column():
                gr.HTML("""
                <div class="stat-card">
                    <h3>📊 Statistics</h3>
                    <p><strong>Total Images:</strong> 100</p>
                    <p><strong>Model:</strong> FLUX.1-Krea-dev</p>
                    <p><strong>Workflow:</strong> ComfyUI</p>
                </div>
                """)
            
            with gr.Column():
                gr.HTML("""
                <div class="stat-card">
                    <h3>⚙️ Technical Details</h3>
                    <p><strong>Log File:</strong> comfyui_8000.log</p>
                    <p><strong>Workflow Screenshot:</strong> Screenshot 2025-10-24 at 11.06.28 AM.png</p>
                    <p><strong>Image Range:</strong> flux_krea_00370_.png to flux_krea_00469_.png</p>
                </div>
                """)
        
        # Main content in tabs
        with gr.Tabs():
            
            # Gallery Tab
            with gr.TabItem("🖼️ Image Gallery"):
                gr.Markdown("### FLUX.1-Krea-dev Generated Images (100 Total)")
                
                with gr.Row():
                    gallery = gr.Gallery(
                        value=create_flux_krea_gallery(),
                        label="FLUX.1-Krea-dev Image Collection",
                        columns=4,
                        rows=4,
                        height="auto",
                        object_fit="cover",
                        show_label=True
                    )
                
                with gr.Row():
                    gr.Button("Refresh Gallery", variant="primary")
                    gr.Button("Export Gallery", variant="secondary")
            
            # Workflow Tab
            with gr.TabItem("🔧 Workflow"):
                gr.Markdown("### ComfyUI Workflow Configuration")
                
                with gr.Row():
                    with gr.Column():
                        gr.Image(
                            value=get_comfyui_workflow_image(),
                            label="Workflow Screenshot",
                            height=500
                        )
                    
                    with gr.Column():
                        gr.Markdown("""
                        #### Workflow Components
                        
                        **1. Model Loading**
                        - Checkpoint: flux1-krea-dev_fp32.safetensors
                        - VAE: ae.safetensors
                        
                        **2. Text Encoding**
                        - CLIP Text Encode (Positive)
                        - ConditioningZeroOut (Conditioning)
                        
                        **3. Sampling**
                        - Sampler Name: Euler
                        - Steps: 20
                        - CFG: 1.0
                        
                        **4. Output**
                        - VAE Decode
                        - Save Image
                        """)
            
            # Log Tab
            with gr.TabItem("📋 Generation Log"):
                gr.Markdown("### ComfyUI Execution Log")
                
                log_content = gr.Textbox(
                    value=load_comfyui_log(),  # This will load the actual log file
                    label="comfyui_8000.log",
                    lines=30,
                    max_lines=100,
                    show_copy_button=True
                )
                
                with gr.Row():
                    load_log_btn = gr.Button("📄 Reload Log", variant="primary")
                    clear_log_btn = gr.Button("🗑️ Clear Log", variant="secondary")
                    download_log_btn = gr.Button("💾 Download Log", variant="secondary")
        
        # Footer
        gr.Markdown("---")
        gr.Markdown(
            """
            <div style="text-align: center; color: #666;">
            <p>FLUX.1-Krea-dev & ComfyUI Showcase Demo | Created with Gradio</p>
            </div>
            """
        )
        
        # Function bindings for advanced demo
        load_log_btn.click(
            fn=load_comfyui_log,
            outputs=log_content
        )
        
        clear_log_btn.click(
            fn=lambda: "",
            outputs=log_content
        )
    
    return demo

# Create and launch the demo
if __name__ == "__main__":
    # Create necessary directories for file structure
    os.makedirs("images", exist_ok=True)
    os.makedirs("workflows", exist_ok=True)
    os.makedirs("logs", exist_ok=True)
    
    # Check if log file exists, if not create it with a placeholder
    log_file_path = "logs/comfyui_8000.log"
    if not os.path.exists(log_file_path):
        with open(log_file_path, "w", encoding="utf-8") as f:
            f.write("# ComfyUI log file will appear here once generation starts\n")
            f.write("# Please ensure your ComfyUI is generating logs to this file\n")
    
    # Launch the demo
    demo = create_advanced_demo()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True,
        show_error=True
    )