import gradio as gr
import os

def create_showcase_app():
    # Define file paths based on your directory listing
    files = {
        "flux_krea_00101.png": "flux_krea_00101_.png",  # Fixed: your file has underscore
        "flux_krea_00102.png": "flux_krea_00102_.png",
        "output.mp4": "output.mp4",
        "output_2.mp4": "output_2.mp4", 
        "output_3.mp4": "output_3.mp4",
        "action_history": "action_history_mario_best_ep3199.txt",
        "training_progress": "training_progress_20251125_145551.png",
        "training_log": "training_output_20251125_104616.log"
    }
    
    # Check which files actually exist
    existing_files = {name: path for name, path in files.items() if os.path.exists(path)}
    print(f"Found {len(existing_files)} files: {list(existing_files.keys())}")
    
    # Simple Blocks without theme for compatibility
    with gr.Blocks(title="AI Training Showcase") as demo:
        gr.Markdown("<h1 style='text-align: center;'>🎮 AI Training Showcase</h1>")
        gr.Markdown("## Model Performance and Training Progress")
        
        # First row: Center image and first video
        with gr.Row():
            with gr.Column(scale=2):
                if "flux_krea_00101.png" in existing_files:
                    gr.Markdown("### 🖼️ Main Result")
                    gr.Image(
                        value=existing_files["flux_krea_00101.png"],
                        label="flux_krea_00101_.png",
                        show_label=True,
                        height=400
                    )
                else:
                    gr.Markdown("### ❌ flux_krea_00101_.png not found")
            
            with gr.Column(scale=1):
                if "output.mp4" in existing_files:
                    gr.Markdown("### 🎥 Video 1")
                    gr.Video(
                        value=existing_files["output.mp4"],
                        label="output.mp4",
                        height=300
                    )
                else:
                    gr.Markdown("### ❌ output.mp4 not found")
        
        # Second row: Second image and second video
        with gr.Row():
            with gr.Column(scale=2):
                if "flux_krea_00102.png" in existing_files:
                    gr.Markdown("### 🖼️ Secondary Result")
                    gr.Image(
                        value=existing_files["flux_krea_00102.png"],
                        label="flux_krea_00102_.png", 
                        show_label=True,
                        height=400
                    )
                else:
                    gr.Markdown("### ❌ flux_krea_00102_.png not found")
            
            with gr.Column(scale=1):
                if "output_2.mp4" in existing_files:
                    gr.Markdown("### 🎥 Video 2")
                    gr.Video(
                        value=existing_files["output_2.mp4"],
                        label="output_2.mp4",
                        height=300
                    )
                else:
                    gr.Markdown("### ❌ output_2.mp4 not found")
        
        # Third row: Training progress and third video
        with gr.Row():
            with gr.Column(scale=2):
                if "training_progress" in existing_files:
                    gr.Markdown("### 📊 Training Progress")
                    gr.Image(
                        value=existing_files["training_progress"],
                        label="training_progress_20251125_145551.png",
                        show_label=True,
                        height=400
                    )
                else:
                    gr.Markdown("### ❌ training_progress.png not found")
            
            with gr.Column(scale=1):
                if "output_3.mp4" in existing_files:
                    gr.Markdown("### 🎥 Video 3")
                    gr.Video(
                        value=existing_files["output_3.mp4"],
                        label="output_3.mp4",
                        height=300
                    )
                else:
                    gr.Markdown("### ❌ output_3.mp4 not found")
        
        # Fourth row: Text files
        with gr.Row():
            if "action_history" in existing_files:
                with gr.Column():
                    gr.Markdown("### 📝 Action History")
                    gr.File(
                        value=existing_files["action_history"],
                        label="action_history_mario_best_ep3199.txt",
                        height=200
                    )
            
            if "training_log" in existing_files:
                with gr.Column():
                    gr.Markdown("### 📋 Training Log")
                    gr.File(
                        value=existing_files["training_log"],
                        label="training_output_20251125_104616.log",
                        height=200
                    )
        
        # File status information
        with gr.Row():
            with gr.Column():
                gr.Markdown("### 📁 File Status")
                status_text = f"""
                Found {len(existing_files)} out of {len(files)} expected files:
                {chr(10).join(['✅ ' + name for name in existing_files.keys()])}
                """
                if len(existing_files) < len(files):
                    missing = set(files.keys()) - set(existing_files.keys())
                    status_text += f"\n\nMissing files:\n{chr(10).join(['❌ ' + name for name in missing])}"
                
                gr.Markdown(status_text)
    
    return demo

if __name__ == "__main__":
    # Create and launch the app
    app = create_showcase_app()
    app.launch(
        server_name="0.0.0.0",  # Allows external access
        server_port=7860,        # Default Gradio port
        share=False              # Set to True to get a public URL
    )