#!/usr/bin/env python3
"""
Agent2Robot - LLM-Agent-Designed Obstacle-Passing Vehicle System
Gradio User Interface Implementation
Track 3: Agentic Demo Showcase
"""

import os
import ssl
import time
import json
import tempfile
from datetime import datetime
from pathlib import Path

# SSL workaround for Gradio issues
try:
    import certifi
    os.environ['SSL_CERT_FILE'] = certifi.where()
except ImportError:
    pass

try:
    ssl._create_default_https_context = ssl._create_unverified_context
except AttributeError:
    pass

# Import Gradio with error handling
GRADIO_AVAILABLE = False
try:
    import gradio as gr
    GRADIO_AVAILABLE = True
    print("✓ Gradio imported successfully")
except Exception as e:
    print(f"⚠ Gradio import failed: {e}")
    exit(1)

# Import backend components
from main_orchestrator import HackathonVehicleDesigner

# Global configuration
MAX_ITERATIONS = 5
designer = HackathonVehicleDesigner()

def ui_function_wrapper(vehicle_type, user_description):
    """
    Main UI wrapper function that yields real-time updates to multiple Gradio components
    Returns tuples in the order: process_log, current_design_specs, progress_bar, 
    results_accordion, final_status, simulation_video, best_design_specs, download_json, 
    performance_summary, llm_rationale
    """
    global designer
    
    # Reset designer for new task
    designer.reset_design_session()
    designer.vehicle_type = vehicle_type.lower()
    designer.user_task_description = user_description
    
    # Initial setup - yield initial states
    yield (
        "🚀 Initializing Agent2Robot system...\n",  # process_log_output
        {},  # current_design_specs_output
        0,   # progress_bar_output
        gr.Accordion(open=False),  # results_accordion - keep closed initially
        "",  # final_status_output
        None,  # simulation_video_output
        {},  # best_design_specs_output
        None,  # download_json_output
        "",  # performance_summary_output
        ""   # llm_rationale_output
    )
    
    # Parse user criteria
    designer.log_process_step("🎯 Analyzing user task and success criteria...")
    criteria = designer.parse_user_task_for_criteria(user_description)
    
    designer.log_process_step(f"📋 Interpreted success criteria:")
    for criterion in criteria:
        designer.log_process_step(f"  • {criterion}")
    
    # Update with criteria interpretation
    current_log = "\n".join(designer.process_log)
    yield (
        current_log,  # process_log_output
        {"interpreted_criteria": criteria},  # current_design_specs_output
        0,  # progress_bar_output
        gr.Accordion(open=False),  # results_accordion
        "",  # final_status_output
        None,  # simulation_video_output
        {},  # best_design_specs_output
        None,  # download_json_output
        "",  # performance_summary_output
        ""   # llm_rationale_output
    )
    
    # Start design process
    designer.log_process_step(f"🚀 Starting {vehicle_type} design process...")
    designer.log_process_step(f"🎯 Target: {user_description}")
    
    current_log = "\n".join(designer.process_log)
    yield (
        current_log,  # process_log_output
        {"status": "Design process starting..."},  # current_design_specs_output
        0,  # progress_bar_output
        gr.Accordion(open=False),  # results_accordion
        "",  # final_status_output
        None,  # simulation_video_output
        {},  # best_design_specs_output
        None,  # download_json_output
        "",  # performance_summary_output
        ""   # llm_rationale_output
    )
    
    # Run iterations
    for iteration in range(1, MAX_ITERATIONS + 1):
        designer.log_process_step(f"\n=== Starting Iteration {iteration}/{MAX_ITERATIONS} ===")
        
        # Update progress at start of iteration
        current_log = "\n".join(designer.process_log)
        progress_value = (iteration - 0.5) / MAX_ITERATIONS * 100  # Convert to percentage
        yield (
            current_log,  # process_log_output
            {"current_iteration": iteration, "max_iterations": MAX_ITERATIONS, "status": "Running..."},  # current_design_specs_output
            progress_value,  # progress_bar_output
            gr.Accordion(open=False),  # results_accordion
            "",  # final_status_output
            None,  # simulation_video_output
            {},  # best_design_specs_output
            None,  # download_json_output
            "",  # performance_summary_output
            ""   # llm_rationale_output
        )
        
        # Run the iteration
        try:
            success = designer.run_single_iteration(iteration)
            
            # Get current design specs for display
            if designer.all_attempts:
                current_attempt = designer.all_attempts[-1]
                current_specs = current_attempt['vehicle_specs']
                design_reasoning = current_attempt.get('design_reasoning', 'No reasoning provided')
                
                # Update with current iteration results
                current_log = "\n".join(designer.process_log)
                progress_value = iteration / MAX_ITERATIONS * 100
                
                current_specs_display = {
                    "iteration": iteration,
                    "vehicle_specs": current_specs,
                    "design_reasoning_preview": design_reasoning[:200] + "..." if len(design_reasoning) > 200 else design_reasoning,
                    "status": "✅ SUCCESS" if success else "🔄 Completed - Evaluating..."
                }
                
                yield (
                    current_log,  # process_log_output
                    current_specs_display,  # current_design_specs_output
                    progress_value,  # progress_bar_output
                    gr.Accordion(open=False),  # results_accordion
                    "",  # final_status_output
                    None,  # simulation_video_output
                    {},  # best_design_specs_output
                    None,  # download_json_output
                    "",  # performance_summary_output
                    ""   # llm_rationale_output
                )
            
            if success:
                designer.log_process_step("🎉 SUCCESS! Design meets all criteria!")
                break
                
        except Exception as e:
            designer.log_process_step(f"❌ Error in iteration {iteration}: {str(e)}")
            current_log = "\n".join(designer.process_log)
            progress_value = iteration / MAX_ITERATIONS * 100
            yield (
                current_log,  # process_log_output
                {"error": f"Iteration {iteration} failed", "details": str(e)},  # current_design_specs_output
                progress_value,  # progress_bar_output
                gr.Accordion(open=False),  # results_accordion
                "",  # final_status_output
                None,  # simulation_video_output
                {},  # best_design_specs_output
                None,  # download_json_output
                "",  # performance_summary_output
                ""   # llm_rationale_output
            )
    
    # Generate final results
    designer.log_process_step("📊 Generating final results and visualizations...")
    current_log = "\n".join(designer.process_log)
    yield (
        current_log,  # process_log_output
        {"status": "Generating final results..."},  # current_design_specs_output
        100,  # progress_bar_output - complete
        gr.Accordion(open=False),  # results_accordion
        "",  # final_status_output
        None,  # simulation_video_output
        {},  # best_design_specs_output
        None,  # download_json_output
        "",  # performance_summary_output
        ""   # llm_rationale_output
    )
    
    # Prepare final outputs
    if designer.overall_success:
        final_status = "## 🎉 SUCCESS!\n\nThe LLM agent successfully designed a vehicle that meets all criteria!"
        status_emoji = "✅"
    else:
        final_status = "## ⚠️ PROCESS COMPLETED\n\nThe agent completed all iterations. Showing best attempt found."
        status_emoji = "🔄"
    
    # Get best design specs
    best_specs = designer.best_attempt['vehicle_specs'] if designer.best_attempt else {}
    
    # Create visualization
    simulation_gif_path = None
    try:
        simulation_gif_path = designer.create_final_visualization()
    except Exception as e:
        designer.log_process_step(f"⚠️ Error creating visualization: {str(e)}")
    
    # Format performance summary
    if designer.best_attempt:
        eval_results = designer.best_attempt['evaluation_results']
        performance_summary = f"""## 📊 Performance Summary of Best Design

**Iteration Found**: {designer.best_iteration}/{len(designer.all_attempts)}
**Final Position**: {eval_results.get('final_robot_x_position', 0.0):.3f}m
**Crossed Obstacle**: {'✅ Yes' if eval_results.get('robot_crossed_obstacle', False) else '❌ No'}
**Remained Stable**: {'✅ Yes' if eval_results.get('robot_remains_upright', False) else '❌ No'}
**Clean Pass**: {'✅ Yes' if eval_results.get('no_significant_collision_with_obstacle_during_pass', False) else '❌ No'}

**Overall Success**: {'✅ ACHIEVED' if eval_results.get('overall_success', False) else '❌ NOT FULLY ACHIEVED'}

**Target Distance**: 0.8m (obstacle clearance)
**Achieved Distance**: {eval_results.get('final_robot_x_position', 0.0):.3f}m
**Success Rate**: {100 if eval_results.get('overall_success', False) else 0}%

{status_emoji} **Status**: {'Complete Success' if designer.overall_success else 'Best Effort'}
"""
    else:
        performance_summary = "## ❌ No successful attempts recorded\n\nThe system was unable to generate valid designs."
    
    # Get LLM rationale
    llm_rationale = designer.best_attempt['design_reasoning'] if designer.best_attempt else "No design reasoning available"
    
    # Create downloadable specs
    download_specs_path = None
    try:
        download_specs_path = designer.save_design_specs_json()
    except Exception as e:
        designer.log_process_step(f"⚠️ Error saving specs: {str(e)}")
    
    # Final log update
    designer.log_process_step(f"\n🏁 DESIGN PROCESS COMPLETED")
    designer.log_process_step(f"📊 Total iterations: {len(designer.all_attempts)}")
    designer.log_process_step(f"🏆 Best iteration: {designer.best_iteration}")
    designer.log_process_step(f"✅ Overall success: {designer.overall_success}")
    
    final_log = "\n".join(designer.process_log)
    
    # Final yield with all results
    yield (
        final_log,  # process_log_output
        {"final_summary": f"Process completed. {len(designer.all_attempts)} iterations run."},  # current_design_specs_output
        100,  # progress_bar_output
        gr.Accordion(open=True),  # results_accordion - open results section
        final_status,  # final_status_output
        simulation_gif_path,  # simulation_video_output
        best_specs,  # best_design_specs_output
        download_specs_path,  # download_json_output
        performance_summary,  # performance_summary_output
        llm_rationale  # llm_rationale_output
    )

def create_agent2robot_interface():
    """Create the Agent2Robot Gradio interface"""
    
    # Custom CSS for better appearance
    custom_css = """
    .main-header {
        text-align: center;
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
        padding: 30px;
        border-radius: 15px;
        margin-bottom: 20px;
        box-shadow: 0 8px 16px rgba(0,0,0,0.1);
    }
    .process-log {
        font-family: 'Courier New', monospace;
        font-size: 12px;
        line-height: 1.4;
    }
    .success-indicator {
        background: linear-gradient(90deg, #4CAF50, #45a049);
        color: white;
        padding: 10px;
        border-radius: 8px;
        margin: 5px 0;
    }
    .iteration-info {
        background: linear-gradient(90deg, #2196F3, #1976D2);
        color: white;
        padding: 8px;
        border-radius: 6px;
        margin: 3px 0;
    }
    """
    
    with gr.Blocks(
        title="🤖🚁 Agent2Robot - LLM Vehicle Designer",
        theme=gr.themes.Soft(),
        css=custom_css
    ) as demo:
        
        # Header Section
        gr.HTML("""
        <div class="main-header">
            <h1>🤖🚁 Agent2Robot</h1>
            <h2>LLM-Agent-Designed Obstacle-Passing Vehicle System</h2>
            <p><strong>Hackathon Submission - Track 3: Agentic Demo Showcase</strong></p>
            <p>Describe your desired vehicle and task in natural language, then watch our AI agent design, simulate, and optimize it in real-time!</p>
        </div>
        """)
        
        # Main Input Section
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("## 🎯 1. Define Your Vehicle Challenge")
                
                vehicle_type_input = gr.Radio(
                    choices=["Robot", "Drone"], 
                    label="1. Choose Vehicle Type",
                    value="Robot",
                    info="Select whether you want a ground robot or flying drone"
                )
                
                user_description_input = gr.Textbox(
                    lines=5, 
                    label="2. Describe Vehicle's Task & Success Criteria",
                    placeholder="e.g., 'Design a robot that can cross the 5cm box obstacle quickly and without tipping over, then stop safely.' or 'Create a drone that flies over the wall, lands gently 1 meter beyond it, and remains stable.'",
                    value="Design a robot that can cross the 5cm high obstacle smoothly and come to a controlled stop."
                )
                
                start_button = gr.Button(
                    "🚀 Start AI Design Process", 
                    variant="primary",
                    size="lg"
                )
                
                gr.Markdown("""
                ### 📋 Environment Info
                - **Obstacle**: 5cm high × 50cm wide box
                - **Success Target**: Vehicle reaches x > 0.8m
                - **Physics**: Real-time PyBullet simulation
                - **Max Iterations**: 5 design attempts
                """)
            
            with gr.Column(scale=2):
                gr.Markdown("## 🤖 2. Watch the AI Agent Work")
                
                process_log_output = gr.Textbox(
                    label="🤖 AI Agent - Live Process Log",
                    lines=15,
                    interactive=False,
                    show_copy_button=True,
                    elem_classes=["process-log"],
                    placeholder="Process log will appear here in real-time as the AI agent works..."
                )
                
                with gr.Row():
                    current_design_specs_output = gr.JSON(
                        label="⚙️ Current Design Specs Being Tested",
                        interactive=False
                    )
                    
                    progress_bar_output = gr.Slider(
                        minimum=0, 
                        maximum=100,
                        step=1,
                        label="Progress (%)",
                        interactive=False,
                        show_label=True
                    )
        
        # Results Section
        with gr.Accordion("🏆 Final Results & Design Specifications", open=False) as results_accordion:
            final_status_output = gr.Markdown(
                label="🏁 Final Run Status",
                value="Waiting for process to complete..."
            )
            
            with gr.Row():
                with gr.Column(scale=2):
                    simulation_video_output = gr.Image(
                        label="🎬 Simulation of Best Design's Trial",
                        interactive=False,
                        height=300
                    )
                    
                    performance_summary_output = gr.Markdown(
                        label="📊 Performance Summary of Best Design"
                    )
                
                with gr.Column(scale=1):
                    best_design_specs_output = gr.JSON(
                        label="🔩 Best Vehicle Design Specifications",
                        show_label=True
                    )
                    
                    download_json_output = gr.File(
                        label="📄 Download Best Design Specs (JSON)",
                        file_count="single",
                        type="filepath",
                        interactive=True
                    )
                    
                    llm_rationale_output = gr.Textbox(
                        label="💡 LLM's Design Rationale",
                        lines=6,
                        interactive=False,
                        show_copy_button=True
                    )
        
        # Connect button to the wrapper function
        start_button.click(
            fn=ui_function_wrapper,
            inputs=[vehicle_type_input, user_description_input],
            outputs=[
                process_log_output,
                current_design_specs_output,
                progress_bar_output,
                results_accordion,
                final_status_output,
                simulation_video_output,
                best_design_specs_output,
                download_json_output,
                performance_summary_output,
                llm_rationale_output
            ],
            show_progress=False  # We handle progress manually
        )
        
        # Information Footer
        gr.Markdown("---")
        gr.Markdown("""
        ## 🔬 How the Agentic AI Works
        
        1. **🎯 Criteria Interpretation**: AI analyzes your natural language task and defines measurable success conditions
        2. **🔧 Intelligent Design**: LLM proposes vehicle specifications based on physics principles and your requirements  
        3. **⚗️ Physics Simulation**: Each design is tested in accurate PyBullet physics simulation with real obstacles
        4. **📊 Performance Analysis**: Results are evaluated against your interpreted criteria with detailed metrics
        5. **🔄 Iterative Learning**: AI uses simulation feedback to refine and improve designs automatically
        6. **🏆 Best Design Selection**: System tracks performance and presents the optimal solution found
        
        **🚀 Innovation**: This demonstrates autonomous AI that goes beyond text generation - it's an agent that designs, tests, learns, and optimizes physical systems to meet user-defined functional requirements.
        """)
    
    return demo

if __name__ == "__main__":
    print("🤖🚁 Agent2Robot - LLM-Agent-Designed Vehicle System")
    print("=" * 60)
    print("🚀 Launching enhanced Gradio interface...")
    
    try:
        # Create and launch the interface
        app = create_agent2robot_interface()
        app.launch(
            server_name="0.0.0.0",
            server_port=7860,
            share=False,  # Set to True for public sharing
            show_error=True,
            inbrowser=True,
            quiet=False
        )
    except Exception as e:
        print(f"❌ Error launching app: {e}")
        print("Please check your installation and try again.")