"""
CUDA Quantum Python Bridge - Production Version
Provides a unified interface for Node.js to interact with CUDA Quantum functionality
"""

import json
import sys
import traceback
import numpy as np
from typing import Dict, List, Any, Optional, Union, Tuple
import logging

# Setup logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

try:
    import cudaq
    from cudaq import spin
    CUDAQ_AVAILABLE = True
    print("INFO - CUDA Quantum successfully imported", flush=True)
except ImportError as e:
    CUDAQ_AVAILABLE = False
    print(f"WARNING - CUDA Quantum not available: {e}", flush=True)
    print("WARNING - Running in development mode without CUDA Quantum", flush=True)
    
    # Create mock cudaq module for development
    class MockCudaq:
        @staticmethod
        def kernel(func):
            return func
        
        @staticmethod
        def get_targets():
            return ['mock-cpu', 'mock-gpu']
        
        @staticmethod
        def set_target(target, **kwargs):
            pass
        
        @staticmethod
        def sample(kernel, *args, **kwargs):
            # Mock sampling results
            return {'00': 500, '11': 500}
        
        @staticmethod
        def observe(kernel, hamiltonian, *args, **kwargs):
            # Mock observation result
            class MockResult:
                def expectation(self):
                    return 0.0
            return MockResult()
        
        @staticmethod
        def get_state(kernel, *args):
            # Mock state vector
            return [1.0+0j, 0.0+0j]
        
        @staticmethod
        def qview():
            pass
    
    class MockSpin:
        @staticmethod
        def z(index):
            return f"Z{index}"
    
    cudaq = MockCudaq()
    spin = MockSpin()


class QuantumKernelManager:
    """Manages quantum kernels and their execution"""
    
    def __init__(self):
        self.kernels: Dict[str, Any] = {}
        self.kernel_metadata: Dict[str, Dict] = {}
        
    def create_kernel(self, name: str, num_qubits: int, parameters: Optional[List[Dict]] = None) -> Dict:
        """Create a new quantum kernel"""
        if not CUDAQ_AVAILABLE:
            # Mock mode - store kernel metadata
            self.kernel_metadata[name] = {
                "num_qubits": num_qubits,
                "parameters": parameters or [],
                "operations": []
            }
            return {"success": True, "kernel_name": name, "mode": "mock"}
            
        try:
            # Create kernel dynamically
            kernel_code = self._generate_kernel_code(name, num_qubits, parameters or [])
            exec(kernel_code, globals())
            
            self.kernels[name] = globals()[name]
            self.kernel_metadata[name] = {
                "num_qubits": num_qubits,
                "parameters": parameters or [],
                "operations": []
            }
            
            return {"success": True, "kernel_name": name}
            
        except Exception as e:
            logger.error(f"Error creating kernel {name}: {e}")
            return {"error": str(e), "traceback": traceback.format_exc()}
    
    def _generate_kernel_code(self, name: str, num_qubits: int, parameters: List[Dict]) -> str:
        """Generate CUDA Quantum kernel code dynamically"""
        param_str = ""
        if parameters:
            param_types = {
                "float": "float",
                "int": "int", 
                "complex": "complex",
                "angle": "float"
            }
            
            params = []
            for param in parameters:
                param_type = param_types.get(param.get("type", "float"), "float")
                params.append(f"{param['name']}: {param_type}")
            param_str = ", " + ", ".join(params)
        
        kernel_code = f"""
@cudaq.kernel
def {name}(qubits: cudaq.qview{param_str}):
    pass
"""
        return kernel_code
    
    def apply_gate(self, kernel_name: str, gate_name: str, 
                   target_qubits: List[int], control_qubits: Optional[List[int]] = None,
                   parameters: Optional[List[float]] = None, 
                   adjoint: bool = False) -> Dict:
        """Apply a quantum gate to a kernel"""
        if kernel_name not in self.kernel_metadata:
            return {"error": f"Kernel {kernel_name} not found"}
        
        try:
            operation = {
                "gate": gate_name,
                "targets": target_qubits,
                "controls": control_qubits or [],
                "parameters": parameters or [],
                "adjoint": adjoint
            }
            
            self.kernel_metadata[kernel_name]["operations"].append(operation)
            
            return {"success": True, "operation": operation}
            
        except Exception as e:
            logger.error(f"Error applying gate {gate_name} to kernel {kernel_name}: {e}")
            return {"error": str(e), "traceback": traceback.format_exc()}


class QuantumExecutor:
    """Handles quantum circuit execution and measurement"""
    
    def __init__(self):
        self.current_target = None
        
    def set_target(self, target_name: str, **kwargs) -> Dict:
        """Set the quantum computing target/backend"""
        if not CUDAQ_AVAILABLE:
            return {"success": True, "target": target_name, "mode": "mock"}
            
        try:
            # Get available targets first
            available_targets = cudaq.get_targets()
            available_target_names = [str(target) if hasattr(target, '__str__') else repr(target) for target in available_targets]
            
            # Validate target exists
            if target_name not in available_target_names:
                # Try common fallbacks
                fallback_targets = ['qpp-cpu', 'density-matrix-cpu', 'default']
                for fallback in fallback_targets:
                    if fallback in available_target_names:
                        print(f"WARNING - Target {target_name} not available, using {fallback}", flush=True)
                        cudaq.set_target(fallback, **kwargs)
                        self.current_target = fallback
                        return {"success": True, "target": fallback, "fallback": True}
                
                return {"error": f"Invalid target name ({target_name})", "available_targets": available_target_names}
            
            # Set the target
            cudaq.set_target(target_name, **kwargs)
            self.current_target = target_name
            
            return {"success": True, "target": target_name}
            
        except Exception as e:
            try:
                cudaq.reset_target()
            except:
                pass
            logger.error(f"Error setting target {target_name}: {e}")
            return {"error": str(e), "traceback": traceback.format_exc()}
    
    def sample(self, kernel_name: str, shots: int = 1000, 
               parameters: Optional[Dict] = None) -> Dict:
        """Execute quantum circuit and sample measurement results"""
        if not CUDAQ_AVAILABLE:
            # Mock sampling results
            return {
                "success": True,
                "counts": {"00": shots//2, "11": shots//2},
                "shots": shots,
                "most_probable": "00",
                "mode": "mock"
            }
            
        if kernel_name not in kernel_manager.kernels:
            return {"error": f"Kernel {kernel_name} not found"}
        
        try:
            kernel = kernel_manager.kernels[kernel_name]
            
            # Prepare parameters
            args = []
            if parameters:
                # Add parameters in the order they were defined
                metadata = kernel_manager.kernel_metadata[kernel_name]
                for param in metadata.get("parameters", []):
                    if param["name"] in parameters:
                        args.append(parameters[param["name"]])
            
            # Execute with sampling
            if args:
                result = cudaq.sample(kernel, *args, shots_count=shots)
            else:
                result = cudaq.sample(kernel, shots_count=shots)
            
            # Convert results to serializable format
            counts = {}
            for bitstring, count in result.items():
                counts[str(bitstring)] = int(count)
            
            return {
                "success": True,
                "counts": counts,
                "shots": shots,
                "most_probable": result.most_probable()
            }
            
        except Exception as e:
            logger.error(f"Error sampling kernel {kernel_name}: {e}")
            return {"error": str(e), "traceback": traceback.format_exc()}
    
    def observe(self, kernel_name: str, hamiltonian_terms: List[Dict], 
                shots: int = 1000, parameters: Optional[Dict] = None) -> Dict:
        """Calculate expectation value of Hamiltonian"""
        if not CUDAQ_AVAILABLE:
            # Mock observation result
            return {
                "success": True,
                "expectation_value": 0.0,
                "shots": shots,
                "mode": "mock"
            }
            
        if kernel_name not in kernel_manager.kernels:
            return {"error": f"Kernel {kernel_name} not found"}
        
        try:
            kernel = kernel_manager.kernels[kernel_name]
            
            # Build Hamiltonian from terms
            hamiltonian = 0.0
            for term in hamiltonian_terms:
                coeff = term.get("coefficient", 1.0)
                pauli_string = term.get("pauli_string", "")
                
                if pauli_string:
                    hamiltonian += coeff * spin.z(0)  # Simplified - should parse pauli_string
            
            # Prepare parameters
            args = []
            if parameters:
                metadata = kernel_manager.kernel_metadata[kernel_name]
                for param in metadata.get("parameters", []):
                    if param["name"] in parameters:
                        args.append(parameters[param["name"]])
            
            # Calculate expectation value
            if args:
                expectation = cudaq.observe(kernel, hamiltonian, *args, shots_count=shots)
            else:
                expectation = cudaq.observe(kernel, hamiltonian, shots_count=shots)
            
            return {
                "success": True,
                "expectation_value": float(expectation.expectation()),
                "shots": shots
            }
            
        except Exception as e:
            logger.error(f"Error observing kernel {kernel_name}: {e}")
            return {"error": str(e), "traceback": traceback.format_exc()}
    
    def get_state(self, kernel_name: str, parameters: Optional[Dict] = None) -> Dict:
        """Get quantum state vector"""
        if not CUDAQ_AVAILABLE:
            # Mock state vector
            return {
                "success": True,
                "state_vector": [
                    {"real": 1.0, "imag": 0.0},
                    {"real": 0.0, "imag": 0.0}
                ],
                "num_qubits": 1,
                "mode": "mock"
            }
            
        if kernel_name not in kernel_manager.kernels:
            return {"error": f"Kernel {kernel_name} not found"}
        
        try:
            kernel = kernel_manager.kernels[kernel_name]
            
            # Prepare parameters
            args = []
            if parameters:
                metadata = kernel_manager.kernel_metadata[kernel_name]
                for param in metadata.get("parameters", []):
                    if param["name"] in parameters:
                        args.append(parameters[param["name"]])
            
            # Get state vector
            if args:
                state = cudaq.get_state(kernel, *args)
            else:
                state = cudaq.get_state(kernel)
            
            # Convert to serializable format
            state_vector = []
            for amplitude in state:
                state_vector.append({
                    "real": float(amplitude.real),
                    "imag": float(amplitude.imag)
                })
            
            return {
                "success": True,
                "state_vector": state_vector,
                "num_qubits": len(state).bit_length() - 1
            }
            
        except Exception as e:
            logger.error(f"Error getting state for kernel {kernel_name}: {e}")
            return {"error": str(e), "traceback": traceback.format_exc()}
    
    def get_available_targets(self) -> Dict:
        """Get list of available quantum targets"""
        try:
            targets = cudaq.get_targets()
            return {"success": True, "targets": list(targets)}
        except Exception as e:
            logger.error(f"Error getting available targets: {e}")
            return {"error": str(e), "traceback": traceback.format_exc()}
    
    def get_target_info(self) -> Dict:
        """Get information about current target"""
        try:
            info = {
                "current_target": self.current_target,
                "available_targets": list(cudaq.get_targets())
            }
            return {"success": True, "info": info}
        except Exception as e:
            logger.error(f"Error getting target info: {e}")
            return {"error": str(e), "traceback": traceback.format_exc()}


def get_available_targets_standalone() -> Dict:
    """Standalone function to get available targets"""
    if not CUDAQ_AVAILABLE:
        return {
            "success": True, 
            "targets": ["mock-cpu", "mock-gpu"],
            "mode": "mock"
        }
        
    try:
        targets = cudaq.get_targets()
        # Convert Target objects to strings for JSON serialization
        target_names = [str(target) if hasattr(target, '__str__') else repr(target) for target in targets]
        return {"success": True, "targets": target_names}
    except Exception as e:
        logger.error(f"Error getting available targets: {e}")
        return {"error": str(e), "traceback": traceback.format_exc()}


def get_target_info_standalone() -> Dict:
    """Standalone function to get target info"""
    if not CUDAQ_AVAILABLE:
        return {
            "success": True,
            "available_targets": ["mock-cpu", "mock-gpu"],
            "mode": "mock"
        }
        
    try:
        targets = cudaq.get_targets()
        # Convert Target objects to strings for JSON serialization
        target_names = [str(target) if hasattr(target, '__str__') else repr(target) for target in targets]
        return {"success": True, "available_targets": target_names}
    except Exception as e:
        logger.error(f"Error getting target info: {e}")
        return {"error": str(e), "traceback": traceback.format_exc()}


def get_platform_info() -> Dict:
    """Get platform and hardware information"""
    if not CUDAQ_AVAILABLE:
        info = {
            "cuda_quantum_version": "mock-0.8.0",
            "available_targets": ["mock-cpu", "mock-gpu"],
            "python_version": sys.version,
            "platform": sys.platform,
            "mode": "mock"
        }
        return {"success": True, "platform_info": info}
        
    try:
        targets = cudaq.get_targets()
        # Convert Target objects to strings for JSON serialization
        target_names = [str(target) if hasattr(target, '__str__') else repr(target) for target in targets]
        
        info = {
            "cuda_quantum_version": getattr(cudaq, '__version__', '0.8.0'),
            "available_targets": target_names,
            "python_version": sys.version,
            "platform": sys.platform
        }
        return {"success": True, "platform_info": info}
    except Exception as e:
        logger.error(f"Error getting platform info: {e}")
        return {"error": str(e), "traceback": traceback.format_exc()}


# Global instances
kernel_manager = QuantumKernelManager()
executor = QuantumExecutor()


def dispatch_command(command: str, **kwargs) -> Dict:
    """Main dispatch function for Python bridge commands"""
    try:
        if command == "create_kernel":
            return kernel_manager.create_kernel(
                kwargs["name"], 
                kwargs["num_qubits"], 
                kwargs.get("parameters")
            )
        
        elif command == "apply_gate":
            return kernel_manager.apply_gate(
                kwargs["kernel_name"],
                kwargs["gate_name"],
                kwargs["target_qubits"],
                kwargs.get("control_qubits"),
                kwargs.get("parameters"),
                kwargs.get("adjoint", False)
            )
        
        elif command == "set_target":
            target_name = kwargs.get("target")
            if target_name is None:
                return {"error": "Missing required parameter 'target'", "available_params": list(kwargs.keys())}
            return executor.set_target(
                target_name,
                **kwargs.get("configuration", {})
            )
        
        elif command == "sample":
            return executor.sample(
                kwargs["kernel_name"],
                kwargs.get("shots", 1000),
                kwargs.get("parameters")
            )
        
        elif command == "observe":
            return executor.observe(
                kwargs["kernel_name"],
                kwargs["hamiltonian_terms"],
                kwargs.get("shots", 1000),
                kwargs.get("parameters")
            )
        
        elif command == "get_state":
            return executor.get_state(
                kwargs["kernel_name"],
                kwargs.get("parameters")
            )
        
        elif command == "get_available_targets":
            return get_available_targets_standalone()
        
        elif command == "get_target_info":
            return get_target_info_standalone()
        
        elif command == "list_kernels":
            return {
                "success": True,
                "kernels": list(kernel_manager.kernels.keys()),
                "metadata": kernel_manager.kernel_metadata
            }
        
        elif command == "get_platform_info":
            return get_platform_info()
        
        else:
            return {"error": f"Unknown command: {command}"}
    
    except Exception as e:
        logger.error(f"Error executing command {command}: {e}")
        return {"error": str(e), "traceback": traceback.format_exc()}


def main():
    """Main event loop for Python bridge"""
    print("CUDA Quantum Python Bridge - Ready", flush=True)
    
    for line in sys.stdin:
        try:
            line = line.strip()
            if not line:
                continue
                
            request = json.loads(line)
            command = request.get("command")
            data = request.get("data", {})
            request_id = request.get("requestId", "")
            
            # Debug logging
            print(f"DEBUG - Processing command: {command} with ID: {request_id}", file=sys.stderr, flush=True)
            
            if command:
                result = dispatch_command(command, **data)
                
                # Check if the result already has success/error structure
                if isinstance(result, dict) and ("success" in result or "error" in result):
                    # Result already has proper structure, just add request ID
                    response = {**result, "requestId": request_id}
                else:
                    # Wrap result in standard response structure
                    response = {
                        "success": True,
                        "data": result,
                        "requestId": request_id
                    }
                print(f"DEBUG - Command {command} completed successfully", file=sys.stderr, flush=True)
            else:
                response = {
                    "success": False,
                    "error": "No command specified",
                    "requestId": request_id
                }
            
            # Ensure response is sent on a single line with explicit flush
            output = json.dumps(response)
            print(output, flush=True)
            sys.stdout.flush()  # Extra flush to ensure it's sent
            
        except Exception as e:
            print(f"ERROR - Exception in main loop: {str(e)}", file=sys.stderr, flush=True)
            print(f"ERROR - Traceback: {traceback.format_exc()}", file=sys.stderr, flush=True)
            error_response = {
                "success": False,
                "error": str(e),
                "traceback": traceback.format_exc(),
                "requestId": request.get("requestId", "") if 'request' in locals() else ""
            }
            output = json.dumps(error_response)
            print(output, flush=True)
            sys.stdout.flush()  # Extra flush to ensure it's sent


if __name__ == "__main__":
    main()