mcp_manager.py

"""
MCP Manager for NoSQL2SQL
Manages Model Context Protocol servers and AI integration
"""

import os
import sys
import json
import asyncio
import subprocess
from pathlib import Path
from typing import Optional, Dict, Any, List
import google.generativeai as genai


class MCPManager:
    """Manages MCP servers and AI interactions"""
    
    def __init__(self, config_manager):
        self.config = config_manager
        self.mcp_servers = {}
        self.gemini_model = None
        self.is_initialized = False
        
        # MCP server processes
        self.neo4j_server_process = None
        self.redis_server_process = None
        self.couchdb_server_process = None
        
    def initialize(self):
        """Initialize MCP system and AI model"""
        if not self.config.is_mcp_enabled():
            print("MCP is disabled")
            return False
        
        # Initialize Gemini API
        api_key = self.config.get_secure("gemini_api_key")
        if not api_key:
            print("Warning: Gemini API key not configured")
            return False
        
        try:
            genai.configure(api_key=api_key)
            self.gemini_model = genai.GenerativeModel('gemini-2.0-flash-exp')
            print("✅ Gemini 2.0 Flash initialized")
            self.is_initialized = True
            return True
        except Exception as e:
            print(f"❌ Failed to initialize Gemini: {e}")
            return False
    
    def start_mcp_servers(self):
        """Start all MCP servers based on configuration"""
        if not self.config.is_mcp_enabled():
            return False
        
        success = True
        
        # Start Neo4j MCP server
        if not self.start_neo4j_server():
            print("Warning: Could not start Neo4j MCP server")
            success = False
        
        # Start Redis MCP server
        if not self.start_redis_server():
            print("Warning: Could not start Redis MCP server")
            success = False
        
        # Start CouchDB MCP server
        if not self.start_couchdb_server():
            print("Warning: Could not start CouchDB MCP server")
            success = False
        
        return success
    
    def stop_mcp_servers(self):
        """Stop all running MCP servers"""
        servers = [
            ("Neo4j", self.neo4j_server_process),
            ("Redis", self.redis_server_process),
            ("CouchDB", self.couchdb_server_process)
        ]
        
        for name, process in servers:
            if process and process.poll() is None:
                try:
                    process.terminate()
                    process.wait(timeout=5)
                    print(f"✅ Stopped {name} MCP server")
                except:
                    process.kill()
                    print(f"⚠️  Force killed {name} MCP server")
    
    def start_neo4j_server(self):
        """Start Neo4j MCP server"""
        try:
            # Check if mcp_neo4j.py exists in mcp_servers folder
            neo4j_mcp_path = Path(__file__).parent / "mcp_servers" / "mcp_neo4j.py"
            if not neo4j_mcp_path.exists():
                print(f"Neo4j MCP server not found at {neo4j_mcp_path}")
                return False
            
            # Start the server process
            self.neo4j_server_process = subprocess.Popen(
                [sys.executable, str(neo4j_mcp_path)],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                creationflags=subprocess.CREATE_NO_WINDOW if os.name == 'nt' else 0
            )
            
            print("✅ Started Neo4j MCP server")
            return True
        except Exception as e:
            print(f"❌ Failed to start Neo4j MCP server: {e}")
            return False
    
    def start_redis_server(self):
        """Start Redis MCP server"""
        try:
            # Check if mcp_redis.py exists in mcp_servers folder
            redis_mcp_path = Path(__file__).parent / "mcp_servers" / "mcp_redis.py"
            if not redis_mcp_path.exists():
                print(f"Redis MCP server not found at {redis_mcp_path}")
                return False
            
            # Start the server process
            self.redis_server_process = subprocess.Popen(
                [sys.executable, str(redis_mcp_path)],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                creationflags=subprocess.CREATE_NO_WINDOW if os.name == 'nt' else 0
            )
            
            print("✅ Started Redis MCP server")
            return True
        except Exception as e:
            print(f"❌ Failed to start Redis MCP server: {e}")
            return False
    
    def start_couchdb_server(self):
        """Start CouchDB MCP server"""
        try:
            # Check if mcp_couchdb.py exists in mcp_servers folder
            couchdb_mcp_path = Path(__file__).parent / "mcp_servers" / "mcp_couchdb.py"
            if not couchdb_mcp_path.exists():
                print(f"CouchDB MCP server not found at {couchdb_mcp_path}")
                return False
            
            # Start the server process
            self.couchdb_server_process = subprocess.Popen(
                [sys.executable, str(couchdb_mcp_path)],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                creationflags=subprocess.CREATE_NO_WINDOW if os.name == 'nt' else 0
            )
            
            print("✅ Started CouchDB MCP server")
            return True
        except Exception as e:
            print(f"❌ Failed to start CouchDB MCP server: {e}")
            return False
    
    async def query_gemini(self, prompt: str, context: Optional[Dict] = None) -> str:
        """Query Gemini AI model"""
        if not self.is_initialized:
            raise Exception("MCP Manager not initialized")
        
        try:
            # Build full prompt with context
            full_prompt = self._build_prompt(prompt, context)
            
            # Generate response
            response = self.gemini_model.generate_content(full_prompt)
            return response.text
        except Exception as e:
            raise Exception(f"Gemini query failed: {e}")
    
    def query_gemini_sync(self, prompt: str, context: Optional[Dict] = None) -> str:
        """Synchronous wrapper for Gemini query"""
        try:
            # Build full prompt with context
            full_prompt = self._build_prompt(prompt, context)
            
            # Generate response
            response = self.gemini_model.generate_content(full_prompt)
            return response.text
        except Exception as e:
            raise Exception(f"Gemini query failed: {e}")
    
    def _build_prompt(self, prompt: str, context: Optional[Dict] = None) -> str:
        """Build complete prompt with context"""
        if not context:
            return prompt
        
        context_str = "Context:\n"
        for key, value in context.items():
            if isinstance(value, (dict, list)):
                context_str += f"{key}: {json.dumps(value, indent=2)}\n"
            else:
                context_str += f"{key}: {value}\n"
        
        return f"{context_str}\n\nQuery: {prompt}"
    
    def recommend_relationships(self, schema_data: Dict) -> Dict[str, Any]:
        """Use AI to recommend database relationships"""
        if not self.is_initialized:
            return {"error": "MCP not initialized"}
        
        prompt = """
        Analyze the following database schema and recommend relationships between tables.
        For each relationship, provide:
        1. Source table and column
        2. Target table and column
        3. Relationship type (one-to-one, one-to-many, many-to-many)
        4. Confidence score (0-100)
        5. Reasoning
        
        Return the response as a JSON array of relationship objects.
        """
        
        context = {
            "schema": schema_data,
            "task": "relationship_recommendation"
        }
        
        try:
            response = self.query_gemini_sync(prompt, context)
            
            # Parse JSON from response
            # Extract JSON from markdown code blocks if present
            if "```json" in response:
                json_start = response.find("```json") + 7
                json_end = response.find("```", json_start)
                json_str = response[json_start:json_end].strip()
            elif "```" in response:
                json_start = response.find("```") + 3
                json_end = response.find("```", json_start)
                json_str = response[json_start:json_end].strip()
            else:
                json_str = response
            
            recommendations = json.loads(json_str)
            return {
                "success": True,
                "recommendations": recommendations
            }
        except Exception as e:
            return {
                "success": False,
                "error": str(e)
            }
    
    def recommend_relationships_with_context(self, migration_context: Dict) -> Dict[str, Any]:
        """Use AI to recommend database relationships with full migration context"""
        if not self.is_initialized:
            return {"success": False, "error": "MCP not initialized"}
        
        source_db = migration_context.get("source_database", "Unknown")
        target_db = migration_context.get("target_database", "MySQL")
        source_schema = migration_context.get("source_schema", {})
        
        prompt = f"""
        You are an expert database migration consultant helping migrate from {source_db} to {target_db}.
        
        IMPORTANT: You MUST provide at least 3-5 relationship recommendations based on the schema below.
        
        Analyze the source database schema and provide:
        1. Recommended table structure for {target_db}
        2. Foreign key relationships between tables (MINIMUM 3 relationships required)
        3. Data type mappings from {source_db} to {target_db}
        4. Confidence score for each recommendation (0-100)
        5. Clear reasoning for each decision
        
        Source Database: {source_db}
        Schema Details:
        {json.dumps(source_schema, indent=2)}
        
        Return VALID JSON with this EXACT structure:
        {{
          "tables": {{
            "table_name": {{
              "fields": {{
                "id": "INT PRIMARY KEY AUTO_INCREMENT",
                "field_name": "VARCHAR(255)",
                ...
              }}
            }}
          }},
          "recommendations": [
            {{
              "from_table": "posts",
              "from_field": "user_id",
              "to_table": "users",
              "to_field": "id",
              "type": "many-to-one",
              "confidence": 95,
              "reasoning": "Posts belong to users - standard blog relationship"
            }},
            ... (provide at least 3 relationships)
          ]
        }}
        
        Rules:
        - MUST include at least 3-5 relationship recommendations
        - Use proper {target_db} data types
        - Follow database normalization best practices
        - Consider the nature of {source_db} when suggesting relationships
        - For Neo4j: Convert graph relationships to foreign keys
        - For CouchDB/Redis: Infer relationships from document structure and key patterns
        
        Return ONLY valid JSON, no markdown formatting.
        """
        
        try:
            response = self.gemini_model.generate_content(prompt)
            response_text = response.text
            
            # Extract JSON from response
            if "```json" in response_text:
                json_start = response_text.find("```json") + 7
                json_end = response_text.find("```", json_start)
                json_str = response_text[json_start:json_end].strip()
            elif "```" in response_text:
                json_start = response_text.find("```") + 3
                json_end = response_text.find("```", json_start)
                json_str = response_text[json_start:json_end].strip()
            else:
                # Try to find JSON in the response
                start_idx = response_text.find('{')
                end_idx = response_text.rfind('}') + 1
                if start_idx != -1 and end_idx > start_idx:
                    json_str = response_text[start_idx:end_idx]
                else:
                    return {
                        "success": False,
                        "error": "Could not parse JSON from AI response"
                    }
            
            result = json.loads(json_str)
            recommendations = result.get("recommendations", [])
            
            # Ensure at least 1 recommendation
            if not recommendations:
                # Generate a default recommendation based on source DB
                if source_schema.get('database_type') == 'Neo4j':
                    nodes = source_schema.get('nodes', {})
                    node_names = list(nodes.keys())[:2]
                    if len(node_names) >= 2:
                        recommendations = [{
                            "from_table": node_names[0].lower(),
                            "from_field": "id",
                            "to_table": node_names[1].lower(),
                            "to_field": "id",
                            "type": "many-to-many",
                            "confidence": 70,
                            "reasoning": f"AI suggested relationship between {node_names[0]} and {node_names[1]} based on graph structure"
                        }]
                else:
                    recommendations = [{
                        "from_table": "main_table",
                        "from_field": "reference_id",
                        "to_table": "reference_table",
                        "to_field": "id",
                        "type": "many-to-one",
                        "confidence": 60,
                        "reasoning": "AI-suggested basic relationship for normalized schema"
                    }]
            
            return {
                "success": True,
                "recommendations": recommendations,
                "schema_preview": result.get("tables", source_schema)
            }
        except json.JSONDecodeError as e:
            # Fallback: provide default recommendation
            return {
                "success": True,
                "recommendations": [{
                    "from_table": "entity_table",
                    "from_field": "reference_id",
                    "to_table": "reference_table",
                    "to_field": "id",
                    "type": "many-to-one",
                    "confidence": 50,
                    "reasoning": f"AI parsing failed, providing default recommendation. Error: {str(e)}"
                }],
                "schema_preview": source_schema
            }
        except Exception as e:
            return {
                "success": False,
                "error": str(e)
            }
    
    def analyze_schema_with_ai(self, schema_data: Dict) -> Dict[str, Any]:
        """Use AI to analyze and improve schema design"""
        if not self.is_initialized:
            return {"error": "MCP not initialized"}
        
        prompt = """
        Analyze this database schema and provide:
        1. Potential data type improvements
        2. Missing indexes recommendations
        3. Normalization suggestions
        4. Performance optimization tips
        
        Return as JSON with categories: type_improvements, index_recommendations, 
        normalization_suggestions, performance_tips
        """
        
        context = {
            "schema": schema_data,
            "task": "schema_analysis"
        }
        
        try:
            response = self.query_gemini_sync(prompt, context)
            
            # Parse JSON from response
            if "```json" in response:
                json_start = response.find("```json") + 7
                json_end = response.find("```", json_start)
                json_str = response[json_start:json_end].strip()
            elif "```" in response:
                json_start = response.find("```") + 3
                json_end = response.find("```", json_start)
                json_str = response[json_start:json_end].strip()
            else:
                json_str = response
            
            analysis = json.loads(json_str)
            return {
                "success": True,
                "analysis": analysis
            }
        except Exception as e:
            return {
                "success": False,
                "error": str(e)
            }
    
    def chat_query(self, message: str, conversation_history: List[Dict] = None) -> str:
        """Process a chat message with context from conversation history"""
        if not self.is_initialized:
            return "MCP is not initialized. Please enable MCP and configure API key."
        
        try:
            # Check if message contains structured context (starts with ===)
            if "=== CURRENT APPLICATION STATE ===" in message:
                # Message already has full context, use it directly
                full_prompt = f"""You are an AI assistant helping with database migration from NoSQL to MySQL.

You have access to the complete application state and context below. Use this information to provide accurate, helpful responses.

{message}

Based on the context above, provide a helpful response. If schema data is available, analyze it. If recommendations exist, reference them. If migration is complete, confirm the results."""
            else:
                # Build conversation context for simple queries
                context_messages = []
                if conversation_history:
                    for msg in conversation_history[-5:]:  # Last 5 messages for context
                        context_messages.append(f"{msg['role']}: {msg['content']}")
                
                context_str = "\n".join(context_messages) if context_messages else ""
                
                full_prompt = f"""You are an AI assistant helping with database migration from NoSQL to SQL.
                
Previous conversation:
{context_str}

User: {message}

Provide helpful, concise responses about database migration, schema design, and data transformation."""
            
            response = self.gemini_model.generate_content(full_prompt)
            return response.text
        except Exception as e:
            return f"Error: {str(e)}"
    
    def cleanup(self):
        """Cleanup resources"""
        self.stop_mcp_servers()
        self.is_initialized = False


if __name__ == "__main__":
    # Test MCP Manager
    from config_manager import SettingsManager
    
    print("Testing MCP Manager...")
    
    config = SettingsManager()
    config.set("mcp_enabled", True)
    
    # For testing, you would need to set your API key
    # config.set_secure("gemini_api_key", "your-api-key-here")
    
    manager = MCPManager(config)
    
    if manager.initialize():
        print("✅ MCP Manager initialized successfully")
        
        # Test chat
        response = manager.chat_query("What is database normalization?")
        print(f"Chat response: {response[:100]}...")
    else:
        print("❌ MCP Manager initialization failed")
    
    manager.cleanup()