#!/usr/bin/env python3
"""
CUDA Quantum MCP Server - Integration Example
Demonstrates full quantum computing workflow using the MCP server
"""

import json
import subprocess
import sys
import time
from typing import Dict, Any, List

class MCPQuantumClient:
    """Simple client for testing MCP Quantum Server functionality"""
    
    def __init__(self, server_path: str = "dist/index.js"):
        self.server_path = server_path
        self.process = None
        
    def start_server(self):
        """Start the MCP server process"""
        try:
            self.process = subprocess.Popen(
                ["node", self.server_path],
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                bufsize=1
            )
            print("✅ MCP Quantum Server started")
            return True
        except Exception as e:
            print(f"❌ Failed to start server: {e}")
            return False
    
    def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Call an MCP tool and return the response"""
        if not self.process:
            raise RuntimeError("Server not started")
        
        request = {
            "jsonrpc": "2.0",
            "id": int(time.time() * 1000),
            "method": "tools/call",
            "params": {
                "name": tool_name,
                "arguments": arguments
            }
        }
        
        # Send request
        request_json = json.dumps(request) + "\n"
        self.process.stdin.write(request_json)
        self.process.stdin.flush()
        
        # Read response
        response_line = self.process.stdout.readline()
        if response_line:
            try:
                return json.loads(response_line)
            except json.JSONDecodeError:
                return {"error": "Invalid JSON response"}
        
        return {"error": "No response from server"}
    
    def close(self):
        """Close the MCP server process"""
        if self.process:
            self.process.terminate()
            self.process.wait()
            print("🔴 MCP Quantum Server stopped")

def run_quantum_teleportation_example():
    """Run a complete quantum teleportation example"""
    print("\n🚀 Quantum Teleportation Example")
    print("=" * 50)
    
    client = MCPQuantumClient()
    
    try:
        # Start the server
        if not client.start_server():
            return False
        
        time.sleep(2)  # Give server time to initialize
        
        # Step 1: Create quantum teleportation circuit
        print("\n1️⃣  Creating quantum teleportation circuit...")
        
        response = client.call_tool("create_quantum_kernel", {
            "name": "teleportation",
            "num_qubits": 3,
            "description": "Quantum teleportation protocol"
        })
        
        if response.get("error"):
            print(f"❌ Error: {response['error']}")
            return False
        
        print("✅ Quantum kernel created")
        
        # Step 2: Prepare Bell pair (qubits 1,2)
        print("\n2️⃣  Preparing Bell pair...")
        
        # Hadamard on qubit 1
        client.call_tool("apply_quantum_gate", {
            "kernel_name": "teleportation",
            "gate_name": "h",
            "qubits": [1]
        })
        
        # CNOT between qubits 1,2
        client.call_tool("apply_quantum_gate", {
            "kernel_name": "teleportation", 
            "gate_name": "x",
            "qubits": [2],
            "controls": [1]
        })
        
        print("✅ Bell pair prepared")
        
        # Step 3: Teleportation protocol
        print("\n3️⃣  Applying teleportation protocol...")
        
        # CNOT between qubits 0,1
        client.call_tool("apply_quantum_gate", {
            "kernel_name": "teleportation",
            "gate_name": "x", 
            "qubits": [1],
            "controls": [0]
        })
        
        # Hadamard on qubit 0
        client.call_tool("apply_quantum_gate", {
            "kernel_name": "teleportation",
            "gate_name": "h",
            "qubits": [0]
        })
        
        print("✅ Teleportation gates applied")
        
        # Step 4: Sample the circuit
        print("\n4️⃣  Sampling quantum circuit...")
        
        response = client.call_tool("sample_quantum_circuit", {
            "kernel_name": "teleportation",
            "shots": 1000
        })
        
        print("✅ Circuit sampled successfully")
        print("Results:", response.get("result", {}).get("content", [{}])[0].get("text", ""))
        
        # Step 5: Visualize the circuit
        print("\n5️⃣  Visualizing circuit...")
        
        response = client.call_tool("visualize_circuit", {
            "kernel_name": "teleportation"
        })
        
        print("✅ Circuit visualization:")
        print(response.get("result", {}).get("content", [{}])[0].get("text", ""))
        
        return True
        
    except Exception as e:
        print(f"❌ Error in example: {e}")
        return False
    
    finally:
        client.close()

def run_vqe_example():
    """Run a Variational Quantum Eigensolver example"""
    print("\n🧮 VQE H2 Molecule Example") 
    print("=" * 50)
    
    client = MCPQuantumClient()
    
    try:
        if not client.start_server():
            return False
            
        time.sleep(2)
        
        # Step 1: Create parameterized ansatz
        print("\n1️⃣  Creating VQE ansatz...")
        
        client.call_tool("create_quantum_kernel", {
            "name": "h2_ansatz",
            "num_qubits": 4,
            "parameters": [
                {"name": "theta1", "type": "float"},
                {"name": "theta2", "type": "float"}
            ]
        })
        
        # Step 2: Build ansatz circuit
        print("\n2️⃣  Building ansatz circuit...")
        
        # Initial Hartree-Fock state preparation would go here
        # For simplicity, we'll just apply some rotation gates
        
        client.call_tool("apply_quantum_gate", {
            "kernel_name": "h2_ansatz",
            "gate_name": "ry",
            "qubits": [0],
            "parameters": [0.5]  # This would be theta1 in real implementation
        })
        
        client.call_tool("apply_quantum_gate", {
            "kernel_name": "h2_ansatz",
            "gate_name": "ry", 
            "qubits": [1],
            "parameters": [1.2]  # This would be theta2 in real implementation
        })
        
        # Entangling gates
        client.call_tool("apply_quantum_gate", {
            "kernel_name": "h2_ansatz",
            "gate_name": "x",
            "qubits": [1],
            "controls": [0]
        })
        
        print("✅ Ansatz circuit built")
        
        # Step 3: Define H2 Hamiltonian
        print("\n3️⃣  Computing expectation value...")
        
        h2_hamiltonian = [
            {
                "paulis": ["Z", "I", "I", "I"],
                "qubits": [0, 1, 2, 3],
                "coefficient": {"real": -1.0523732, "imag": 0.0}
            },
            {
                "paulis": ["I", "Z", "I", "I"],
                "qubits": [0, 1, 2, 3],
                "coefficient": {"real": -1.0523732, "imag": 0.0}
            },
            {
                "paulis": ["Z", "Z", "I", "I"],
                "qubits": [0, 1, 2, 3],
                "coefficient": {"real": -0.39793742, "imag": 0.0}
            }
        ]
        
        response = client.call_tool("observe_hamiltonian", {
            "kernel_name": "h2_ansatz",
            "hamiltonian_terms": h2_hamiltonian,
            "shots": 10000
        })
        
        print("✅ Expectation value computed")
        print("Results:", response.get("result", {}).get("content", [{}])[0].get("text", ""))
        
        return True
        
    except Exception as e:
        print(f"❌ Error in VQE example: {e}")
        return False
        
    finally:
        client.close()

def run_gpu_acceleration_example():
    """Demonstrate GPU acceleration capabilities"""
    print("\n🎮 GPU Acceleration Example")
    print("=" * 50)
    
    client = MCPQuantumClient()
    
    try:
        if not client.start_server():
            return False
            
        time.sleep(2)
        
        # Step 1: Configure GPU acceleration
        print("\n1️⃣  Configuring GPU acceleration...")
        
        response = client.call_tool("configure_gpu_acceleration", {
            "enable": True,
            "device_id": 0,
            "target": "qpp-gpu"
        })
        
        print("✅ GPU configuration:")
        print(response.get("result", {}).get("content", [{}])[0].get("text", ""))
        
        # Step 2: Create larger circuit for GPU demo
        print("\n2️⃣  Creating large quantum circuit...")
        
        client.call_tool("create_quantum_kernel", {
            "name": "large_circuit",
            "num_qubits": 16,
            "description": "Large circuit for GPU demonstration"
        })
        
        # Step 3: Apply random gates
        print("\n3️⃣  Applying quantum gates...")
        
        # Hadamard on all qubits
        for i in range(16):
            client.call_tool("apply_quantum_gate", {
                "kernel_name": "large_circuit",
                "gate_name": "h",
                "qubits": [i]
            })
        
        # Entangling gates
        for i in range(15):
            client.call_tool("apply_quantum_gate", {
                "kernel_name": "large_circuit",
                "gate_name": "x",
                "qubits": [i + 1],
                "controls": [i]
            })
        
        print("✅ Large circuit constructed")
        
        # Step 4: Sample with GPU acceleration
        print("\n4️⃣  Sampling with GPU acceleration...")
        
        response = client.call_tool("sample_quantum_circuit", {
            "kernel_name": "large_circuit",
            "shots": 5000,
            "target": "qpp-gpu"
        })
        
        print("✅ GPU-accelerated sampling completed")
        print("Results preview:", response.get("result", {}).get("content", [{}])[0].get("text", "")[:200] + "...")
        
        return True
        
    except Exception as e:
        print(f"❌ Error in GPU example: {e}")
        return False
        
    finally:
        client.close()

def main():
    """Run all integration examples"""
    print("🌟 CUDA Quantum MCP Server - Integration Examples")
    print("=" * 60)
    
    examples = [
        ("Quantum Teleportation", run_quantum_teleportation_example),
        ("Variational Quantum Eigensolver", run_vqe_example), 
        ("GPU Acceleration", run_gpu_acceleration_example)
    ]
    
    results = []
    
    for name, example_func in examples:
        print(f"\n🔥 Running {name} Example...")
        try:
            success = example_func()
            results.append((name, success))
            if success:
                print(f"✅ {name} completed successfully!")
            else:
                print(f"❌ {name} failed!")
        except Exception as e:
            print(f"❌ {name} failed with error: {e}")
            results.append((name, False))
    
    # Summary
    print("\n" + "=" * 60)
    print("📊 INTEGRATION RESULTS SUMMARY")
    print("=" * 60)
    
    total_examples = len(results)
    successful = sum(1 for _, success in results if success)
    
    for name, success in results:
        status = "✅ PASS" if success else "❌ FAIL"
        print(f"{status} {name}")
    
    print(f"\n🎯 Success Rate: {successful}/{total_examples} ({successful/total_examples*100:.1f}%)")
    
    if successful == total_examples:
        print("🎉 All integration examples completed successfully!")
        print("🚀 Your CUDA Quantum MCP Server is ready for production!")
    else:
        print("⚠️  Some examples failed. Check the error messages above.")
        print("💡 Make sure CUDA Quantum is installed and the server is built.")
    
    return successful == total_examples

if __name__ == "__main__":
    try:
        success = main()
        sys.exit(0 if success else 1)
    except KeyboardInterrupt:
        print("\n\n🛑 Integration examples interrupted by user")
        sys.exit(1)
    except Exception as e:
        print(f"\n\n💥 Fatal error: {e}")
        sys.exit(1)