#!/usr/bin/env python3
"""
MCP Ecosystem Integration Script

This script integrates all 9 MCP servers into the KGraph-MCP platform:
1. Backs up current data files
2. Installs expanded knowledge graph data
3. Tests the integration
4. Provides rollback if needed

Run with: python integrate_mcp_ecosystem.py
"""

import json
import logging
import os
import shutil
import subprocess
from datetime import datetime

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

class MCPEcosystemIntegrator:
    """Integrates the comprehensive MCP ecosystem into KGraph-MCP."""

    def __init__(self):
        """Initialize the integrator."""
        self.backup_dir = f"backup_mcp_integration_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        self.data_dir = "data"
        self.current_tools_file = "data/initial_tools.json"
        self.current_prompts_file = "data/initial_prompts.json"
        self.expanded_tools_file = "data/initial_tools_expanded.json"
        self.expanded_prompts_file = "data/initial_prompts_expanded.json"

    def create_backup(self) -> bool:
        """Create backup of current data files."""
        try:
            logger.info("Creating backup of current data files...")
            os.makedirs(self.backup_dir, exist_ok=True)

            # Backup current files
            if os.path.exists(self.current_tools_file):
                shutil.copy2(self.current_tools_file, os.path.join(self.backup_dir, "initial_tools.json"))
                logger.info(f"Backed up tools: {self.current_tools_file}")

            if os.path.exists(self.current_prompts_file):
                shutil.copy2(self.current_prompts_file, os.path.join(self.backup_dir, "initial_prompts.json"))
                logger.info(f"Backed up prompts: {self.current_prompts_file}")

            logger.info(f"✅ Backup created: {self.backup_dir}")
            return True

        except Exception as e:
            logger.error(f"❌ Backup failed: {e}")
            return False

    def validate_expanded_files(self) -> bool:
        """Validate the expanded configuration files."""
        try:
            logger.info("Validating expanded configuration files...")

            # Validate tools file
            if not os.path.exists(self.expanded_tools_file):
                logger.error(f"❌ Expanded tools file not found: {self.expanded_tools_file}")
                return False

            with open(self.expanded_tools_file) as f:
                tools_data = json.load(f)

            if not isinstance(tools_data, list) or len(tools_data) == 0:
                logger.error("❌ Invalid tools data structure")
                return False

            logger.info(f"✅ Found {len(tools_data)} tools in expanded configuration")

            # Validate prompts file
            if not os.path.exists(self.expanded_prompts_file):
                logger.error(f"❌ Expanded prompts file not found: {self.expanded_prompts_file}")
                return False

            with open(self.expanded_prompts_file) as f:
                prompts_data = json.load(f)

            if not isinstance(prompts_data, list) or len(prompts_data) == 0:
                logger.error("❌ Invalid prompts data structure")
                return False

            logger.info(f"✅ Found {len(prompts_data)} prompts in expanded configuration")

            # Validate tool-prompt relationships
            tool_ids = {tool["tool_id"] for tool in tools_data}
            prompt_tool_ids = {prompt["target_tool_id"] for prompt in prompts_data}

            missing_tools = prompt_tool_ids - tool_ids
            if missing_tools:
                logger.warning(f"⚠️ Prompts reference missing tools: {missing_tools}")

            logger.info("✅ Configuration files validated successfully")
            return True

        except Exception as e:
            logger.error(f"❌ Validation failed: {e}")
            return False

    def install_expanded_configuration(self) -> bool:
        """Install the expanded configuration files."""
        try:
            logger.info("Installing expanded MCP ecosystem configuration...")

            # Copy expanded files to main locations
            shutil.copy2(self.expanded_tools_file, self.current_tools_file)
            logger.info(f"✅ Installed expanded tools: {self.current_tools_file}")

            shutil.copy2(self.expanded_prompts_file, self.current_prompts_file)
            logger.info(f"✅ Installed expanded prompts: {self.current_prompts_file}")

            logger.info("✅ Expanded configuration installed successfully")
            return True

        except Exception as e:
            logger.error(f"❌ Installation failed: {e}")
            return False

    def test_integration(self) -> bool:
        """Test the integrated ecosystem."""
        try:
            logger.info("Testing integrated MCP ecosystem...")

            # Test loading the configuration
            result = subprocess.run([
                "python", "-c",
                """
import sys
sys.path.append(".")
import app
try:
    # Test agent initialization
    planner, executor = app.initialize_agent_system()
    if planner is None or executor is None:
        print("FAILED: Agent initialization failed")
        sys.exit(1)
    
    # Test knowledge graph loading
    if not hasattr(planner, "kg") or planner.kg is None:
        print("FAILED: Knowledge graph not loaded")
        sys.exit(1)
        
    print(f"SUCCESS: Loaded {len(planner.kg.tools)} tools and {len(planner.kg.prompts)} prompts")
    print("SUCCESS: Integration test passed")
    
except Exception as e:
    print(f"FAILED: {e}")
    sys.exit(1)
                """
            ], capture_output=True, text=True, timeout=60, check=False)

            if result.returncode == 0:
                logger.info("✅ Integration test passed")
                logger.info(result.stdout.strip())
                return True
            logger.error("❌ Integration test failed")
            logger.error(result.stderr.strip())
            return False

        except subprocess.TimeoutExpired:
            logger.error("❌ Integration test timed out")
            return False
        except Exception as e:
            logger.error(f"❌ Integration test error: {e}")
            return False

    def rollback(self) -> bool:
        """Rollback to previous configuration."""
        try:
            logger.info("Rolling back to previous configuration...")

            backup_tools = os.path.join(self.backup_dir, "initial_tools.json")
            backup_prompts = os.path.join(self.backup_dir, "initial_prompts.json")

            if os.path.exists(backup_tools):
                shutil.copy2(backup_tools, self.current_tools_file)
                logger.info("✅ Restored tools configuration")

            if os.path.exists(backup_prompts):
                shutil.copy2(backup_prompts, self.current_prompts_file)
                logger.info("✅ Restored prompts configuration")

            logger.info("✅ Rollback completed successfully")
            return True

        except Exception as e:
            logger.error(f"❌ Rollback failed: {e}")
            return False

    def display_integration_summary(self) -> None:
        """Display summary of the integration."""
        try:
            with open(self.current_tools_file) as f:
                tools = json.load(f)
            with open(self.current_prompts_file) as f:
                prompts = json.load(f)

            print("\n" + "="*60)
            print("🚀 MCP ECOSYSTEM INTEGRATION COMPLETE")
            print("="*60)
            print(f"📊 Total Tools: {len(tools)}")
            print(f"📋 Total Prompts: {len(prompts)}")
            print()

            # Group tools by category
            categories = {}
            for tool in tools:
                tags = tool.get("tags", [])
                category = self._categorize_tool(tags)
                if category not in categories:
                    categories[category] = []
                categories[category].append(tool["name"])

            print("🔧 TOOL CATEGORIES:")
            for category, tool_names in categories.items():
                print(f"  {category}: {len(tool_names)} tools")
                for name in tool_names:
                    print(f"    • {name}")
                print()

            # Display prompt distribution
            prompt_by_tool = {}
            for prompt in prompts:
                tool_id = prompt["target_tool_id"]
                if tool_id not in prompt_by_tool:
                    prompt_by_tool[tool_id] = []
                prompt_by_tool[tool_id].append(prompt["name"])

            print("📋 PROMPT DISTRIBUTION:")
            for tool in tools:
                tool_id = tool["tool_id"]
                prompt_count = len(prompt_by_tool.get(tool_id, []))
                print(f"  {tool['name']}: {prompt_count} specialized prompts")

            print("\n🎯 NEXT STEPS:")
            print("1. Deploy additional MCP servers to HF Spaces")
            print("2. Update main app to showcase 9-tool ecosystem")
            print("3. Test performance across all tools")
            print("4. Create comprehensive demo scenarios")
            print()
            print("✅ Your KGraph-MCP platform now supports 9 comprehensive MCP tools!")

        except Exception as e:
            logger.error(f"❌ Failed to display summary: {e}")

    def _categorize_tool(self, tags: list[str]) -> str:
        """Categorize a tool based on its tags."""
        if any(tag in ["nlp", "text", "summary", "sentiment"] for tag in tags):
            return "📝 Text & NLP"
        if any(tag in ["code", "security", "quality"] for tag in tags):
            return "💻 Code Analysis"
        if any(tag in ["files", "csv", "json", "data"] for tag in tags):
            return "📁 File Processing"
        if any(tag in ["math", "statistics", "calculation"] for tag in tags):
            return "🧮 Mathematics"
        if any(tag in ["web", "scraping", "extraction"] for tag in tags):
            return "🌐 Web Scraping"
        if any(tag in ["vision", "image", "analysis"] for tag in tags):
            return "🖼️ Computer Vision"
        return "🔧 Other Tools"

def main():
    """Main integration process."""
    print("🚀 KGraph-MCP Ecosystem Integration")
    print("="*50)
    print()

    integrator = MCPEcosystemIntegrator()

    try:
        # Step 1: Create backup
        if not integrator.create_backup():
            print("❌ Backup failed - aborting integration")
            return False

        # Step 2: Validate expanded files
        if not integrator.validate_expanded_files():
            print("❌ Configuration validation failed - aborting integration")
            return False

        # Step 3: Install expanded configuration
        if not integrator.install_expanded_configuration():
            print("❌ Installation failed - attempting rollback")
            integrator.rollback()
            return False

        # Step 4: Test integration
        if not integrator.test_integration():
            print("❌ Integration test failed - attempting rollback")
            integrator.rollback()
            return False

        # Step 5: Display summary
        integrator.display_integration_summary()

        print(f"💾 Backup available at: {integrator.backup_dir}")
        print("🎉 Integration completed successfully!")
        return True

    except KeyboardInterrupt:
        print("\n⚠️ Integration interrupted by user")
        print("Rolling back changes...")
        integrator.rollback()
        return False
    except Exception as e:
        logger.error(f"❌ Unexpected error: {e}")
        print("Rolling back changes...")
        integrator.rollback()
        return False

if __name__ == "__main__":
    success = main()
    exit(0 if success else 1)