#!/usr/bin/env python3
"""
Demo script showcasing MCP server integration for LLM-Agent-Designed Vehicle System

This script demonstrates how to interact with our MCP server to design robots and drones
for obstacle crossing using the Model Context Protocol.
"""

import asyncio
import json
import time
from typing import Dict, Any

# Note: In a real implementation, you would use the actual MCP client
# For demo purposes, we'll simulate the MCP calls by importing our server directly
import mcp_server

class MCPDemoClient:
    """Simulated MCP client for demonstration purposes"""
    
    def __init__(self):
        self.server = mcp_server  # Direct import for demo
        
    async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Simulate MCP tool call"""
        print(f"🔧 Calling MCP tool: {tool_name}")
        print(f"📊 Arguments: {json.dumps(arguments, indent=2)}")
        print("─" * 50)
        
        # Call the actual tool function
        tool_func = getattr(self.server, tool_name, None)
        if tool_func:
            result = tool_func(**arguments)
            return result
        else:
            return {"error": f"Tool {tool_name} not found"}

async def demo_obstacle_info():
    """Demo: Get obstacle information"""
    print("🎯 DEMO 1: Getting Obstacle Information")
    print("=" * 60)
    
    client = MCPDemoClient()
    result = await client.call_tool("get_obstacle_info", {})
    
    print("✅ Obstacle Information:")
    print(json.dumps(result, indent=2))
    print("\n")

async def demo_design_parameters():
    """Demo: Get design parameters for robots and drones"""
    print("🔧 DEMO 2: Getting Design Parameters")
    print("=" * 60)
    
    client = MCPDemoClient()
    
    # Get robot parameters
    print("🤖 Robot Design Parameters:")
    robot_params = await client.call_tool("get_design_parameters", {"vehicle_type": "robot"})
    print(json.dumps(robot_params, indent=2))
    print()
    
    # Get drone parameters
    print("🚁 Drone Design Parameters:")
    drone_params = await client.call_tool("get_design_parameters", {"vehicle_type": "drone"})
    print(json.dumps(drone_params, indent=2))
    print("\n")

async def demo_robot_design():
    """Demo: Design and simulate a robot"""
    print("🤖 DEMO 3: Robot Design and Simulation")
    print("=" * 60)
    
    client = MCPDemoClient()
    
    # Design a robot
    robot_specs = {
        "wheel_type": "large_smooth",
        "body_clearance_cm": 7,
        "approach_sensor_enabled": True,
        "main_material": "light_plastic"
    }
    
    design_result = await client.call_tool("design_vehicle", {
        "vehicle_type": "robot",
        "specifications": robot_specs,
        "design_reasoning": "Large smooth wheels for obstacle climbing, moderate clearance for 5cm obstacle, light material for mobility"
    })
    
    print("✅ Robot Design Result:")
    print(json.dumps(design_result, indent=2))
    print()
    
    if design_result.get("success"):
        # Simulate the robot
        print("🏃 Running robot simulation...")
        sim_result = await client.call_tool("simulate_vehicle", {
            "vehicle_type": "robot",
            "specifications": robot_specs,
            "simulation_duration": 5.0
        })
        
        print("✅ Simulation Result:")
        print(json.dumps(sim_result, indent=2))
        print()
        
        if sim_result.get("success"):
            # Evaluate performance
            print("📊 Evaluating robot performance...")
            eval_result = await client.call_tool("evaluate_performance", {
                "simulation_feedback": sim_result["simulation_feedback"]
            })
            
            print("✅ Evaluation Result:")
            print(json.dumps(eval_result, indent=2))
            
            if eval_result.get("success"):
                success_criteria = eval_result["success_criteria"]
                print(f"\n🎯 Performance Summary:")
                print(f"   Crossed Obstacle: {'✅' if success_criteria['crossed_obstacle'] else '❌'}")
                print(f"   Remained Stable: {'✅' if success_criteria['remained_stable'] else '❌'}")
                print(f"   Minimal Collision: {'✅' if success_criteria['minimal_collision'] else '❌'}")
                print(f"   Overall Success: {'✅' if success_criteria['overall_success'] else '❌'}")
    
    print("\n")

async def demo_drone_design():
    """Demo: Design and simulate a drone"""
    print("🚁 DEMO 4: Drone Design and Simulation")
    print("=" * 60)
    
    client = MCPDemoClient()
    
    # Design a drone
    drone_specs = {
        "propeller_size": "medium",
        "flight_height_cm": 25,
        "stability_mode": "auto_hover",
        "main_material": "light_carbon_fiber"
    }
    
    design_result = await client.call_tool("design_vehicle", {
        "vehicle_type": "drone",
        "specifications": drone_specs,
        "design_reasoning": "Medium propellers for balanced thrust, 25cm flight height to safely clear obstacle, auto-hover for stability, carbon fiber for lightweight performance"
    })
    
    print("✅ Drone Design Result:")
    print(json.dumps(design_result, indent=2))
    print()
    
    if design_result.get("success"):
        # Simulate the drone
        print("🛸 Running drone simulation...")
        sim_result = await client.call_tool("simulate_vehicle", {
            "vehicle_type": "drone",
            "specifications": drone_specs,
            "simulation_duration": 5.0
        })
        
        print("✅ Simulation Result:")
        print(json.dumps(sim_result, indent=2))
        print()
        
        if sim_result.get("success"):
            # Evaluate performance
            print("📊 Evaluating drone performance...")
            eval_result = await client.call_tool("evaluate_performance", {
                "simulation_feedback": sim_result["simulation_feedback"]
            })
            
            print("✅ Evaluation Result:")
            print(json.dumps(eval_result, indent=2))
            
            if eval_result.get("success"):
                success_criteria = eval_result["success_criteria"]
                print(f"\n🎯 Performance Summary:")
                print(f"   Crossed Obstacle: {'✅' if success_criteria['crossed_obstacle'] else '❌'}")
                print(f"   Remained Stable: {'✅' if success_criteria['remained_stable'] else '❌'}")
                print(f"   Minimal Collision: {'✅' if success_criteria['minimal_collision'] else '❌'}")
                print(f"   Overall Success: {'✅' if success_criteria['overall_success'] else '❌'}")
    
    print("\n")

async def demo_iterative_design():
    """Demo: Complete iterative design process"""
    print("🔄 DEMO 5: Complete Iterative Design Process")
    print("=" * 60)
    
    client = MCPDemoClient()
    
    # Run iterative design process for a robot
    print("🤖 Running iterative robot design process...")
    iterative_result = await client.call_tool("iterative_design_process", {
        "vehicle_type": "robot",
        "task_description": "Design a robot that can efficiently cross the 5cm obstacle",
        "max_iterations": 3
    })
    
    print("✅ Iterative Design Result:")
    print(json.dumps(iterative_result, indent=2))
    
    if iterative_result.get("success"):
        print(f"\n🎉 SUCCESS! Robot design completed in {iterative_result['iterations_needed']} iterations")
        print(f"🔧 Final Design: {iterative_result['successful_design']}")
    else:
        print(f"\n❌ Design process did not achieve success within max iterations")
        if "best_attempt" in iterative_result:
            print(f"🔧 Best Attempt: {iterative_result['best_attempt'].get('specifications', {})}")
    
    print("\n")

async def demo_comparison():
    """Demo: Compare robot vs drone approaches"""
    print("⚔️ DEMO 6: Robot vs Drone Comparison")
    print("=" * 60)
    
    client = MCPDemoClient()
    
    # Design specifications for comparison
    robot_specs = {
        "wheel_type": "tracked_base",
        "body_clearance_cm": 8,
        "approach_sensor_enabled": True,
        "main_material": "sturdy_metal_alloy"
    }
    
    drone_specs = {
        "propeller_size": "large_stable",
        "flight_height_cm": 20,
        "stability_mode": "auto_hover",
        "main_material": "sturdy_aluminum"
    }
    
    print("🤖 Testing Robust Robot Design...")
    robot_sim = await client.call_tool("simulate_vehicle", {
        "vehicle_type": "robot",
        "specifications": robot_specs,
        "simulation_duration": 8.0
    })
    
    print("🚁 Testing Stable Drone Design...")
    drone_sim = await client.call_tool("simulate_vehicle", {
        "vehicle_type": "drone",
        "specifications": drone_specs,
        "simulation_duration": 8.0
    })
    
    # Evaluate both
    if robot_sim.get("success"):
        robot_eval = await client.call_tool("evaluate_performance", {
            "simulation_feedback": robot_sim["simulation_feedback"]
        })
        
        print("\n🤖 Robot Performance:")
        if robot_eval.get("success"):
            robot_success = robot_eval["success_criteria"]["overall_success"]
            print(f"   Overall Success: {'✅' if robot_success else '❌'}")
            print(f"   Final Position: {robot_eval['evaluation_results'].get('final_robot_x_position', 0):.2f}m")
    
    if drone_sim.get("success"):
        drone_eval = await client.call_tool("evaluate_performance", {
            "simulation_feedback": drone_sim["simulation_feedback"]
        })
        
        print("\n🚁 Drone Performance:")
        if drone_eval.get("success"):
            drone_success = drone_eval["success_criteria"]["overall_success"]
            print(f"   Overall Success: {'✅' if drone_success else '❌'}")
            print(f"   Final Position: {drone_eval['evaluation_results'].get('final_robot_x_position', 0):.2f}m")
    
    print("\n")

async def main():
    """Run all demo scenarios"""
    print("🚀 MCP Vehicle Design System Demo")
    print("=" * 80)
    print("This demo showcases the Model Context Protocol integration")
    print("for our LLM-Agent-Designed Vehicle System.")
    print("=" * 80)
    print()
    
    try:
        # Run all demos
        await demo_obstacle_info()
        await demo_design_parameters()
        await demo_robot_design()
        await demo_drone_design()
        await demo_iterative_design()
        await demo_comparison()
        
        print("🎊 Demo completed successfully!")
        print("\n" + "=" * 80)
        print("💡 MCP Integration Benefits Demonstrated:")
        print("   ✅ Standardized tool access for vehicle design")
        print("   ✅ Support for both robots and drones")
        print("   ✅ Complete design-simulate-evaluate loop")
        print("   ✅ Iterative optimization with LLM feedback")
        print("   ✅ Performance evaluation and comparison")
        print("   ✅ Type-safe parameter validation")
        print("=" * 80)
        
    except Exception as e:
        print(f"❌ Demo failed with error: {e}")
        print("\nNote: This demo requires the enhanced simulation environment.")
        print("Make sure all dependencies are installed and PyBullet is working.")

def run_cli_demo():
    """CLI version that can be run without MCP client"""
    print("🔧 MCP CLI Demo")
    print("=" * 50)
    print("This demonstrates the MCP tools available in our system:")
    print()
    
    # Show available tools
    tools = [
        "get_obstacle_info() - Get obstacle specifications",
        "get_design_parameters(vehicle_type) - Get available parameters",
        "design_vehicle(vehicle_type, specs, reasoning) - Design validation",
        "simulate_vehicle(vehicle_type, specs, duration) - Run simulation",
        "evaluate_performance(feedback) - Evaluate performance",
        "iterative_design_process(vehicle_type, task, max_iter) - Complete process"
    ]
    
    print("📋 Available MCP Tools:")
    for i, tool in enumerate(tools, 1):
        print(f"   {i}. {tool}")
    
    print()
    print("🎯 Example Usage:")
    print("   # Get obstacle info")
    print("   mcp call get_obstacle_info")
    print()
    print("   # Design a robot")
    print("   mcp call design_vehicle \\")
    print("     --vehicle_type robot \\")
    print("     --specifications '{\"wheel_type\": \"large_smooth\", \"body_clearance_cm\": 7}' \\")
    print("     --design_reasoning 'Balanced design for obstacle crossing'")
    print()
    print("   # Run complete design process")
    print("   mcp call iterative_design_process \\")
    print("     --vehicle_type drone \\")
    print("     --max_iterations 3")
    print()
    print("🚀 To use with real MCP client:")
    print("   1. Start server: python mcp_server.py")
    print("   2. Connect client to our MCP server")
    print("   3. Use the tools listed above")

if __name__ == "__main__":
    import sys
    
    if len(sys.argv) > 1 and sys.argv[1] == "--cli":
        run_cli_demo()
    else:
        print("Starting full MCP demo...")
        asyncio.run(main())