# SPDX-License-Identifier: Apache-2.0
# Copyright 2025 SWE Agent
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
MCP (Model Context Protocol) Server Integration for SWE Agent.
Allows users to integrate custom MCP servers with the agent.
"""

import os
import json
import asyncio
import logging
from typing import Dict, List, Any, Optional, Union
from pathlib import Path

try:
    from langchain_mcp_adapters.client import MultiServerMCPClient
    MCP_AVAILABLE = True
except ImportError:
    MCP_AVAILABLE = False

logger = logging.getLogger(__name__)


class MCPIntegration:
    """
    MCP server integration for SWE Agent.
    Manages connection to multiple MCP servers and provides tools to agents.
    """
    
    def __init__(self, config_path: Optional[str] = None):
        """
        Initialize MCP integration.
        
        Args:
            config_path: Path to MCP configuration file (default: .swe_mcp_config.json)
        """
        self.config_path = config_path or ".swe_mcp_config.json"
        self.client: Optional[MultiServerMCPClient] = None
        self.tools: List[Any] = []
        self.servers_config: Dict[str, Dict[str, Any]] = {}
        
        if not MCP_AVAILABLE:
            logger.warning("langchain-mcp-adapters not available. MCP integration disabled.")
            return
            
        self._load_config()
    
    def _load_config(self) -> None:
        """Load MCP server configuration from file."""
        try:
            # Try current directory first
            if Path(self.config_path).exists():
                with open(self.config_path, 'r') as f:
                    self.servers_config = json.load(f)
                logger.info(f"Loaded MCP configuration from {self.config_path}")
                return
            
            # Try package default config as fallback
            package_config = Path(__file__).parent.parent / ".swe_mcp_config.json"
            if package_config.exists():
                with open(package_config, 'r') as f:
                    self.servers_config = json.load(f)
                logger.info(f"Loaded default MCP configuration from package: {package_config}")
                return
                
            logger.info(f"No MCP configuration found at {self.config_path}")
            self._create_default_config()
        except Exception as e:
            logger.error(f"Failed to load MCP configuration: {str(e)}")
            self.servers_config = {}
    
    def _create_default_config(self) -> None:
        """Create default MCP configuration file with DeepWiki and Semgrep servers."""
        default_config = {
            "_description": "Default MCP Server Configuration for SWE Agent PyPI Package",
            "deepwiki": {
                "url": "https://mcp.deepwiki.com/mcp",
                "transport": "streamable_http"
            },
            "semgrep": {
                "url": "https://mcp.semgrep.ai/mcp",
                "transport": "streamable_http"
            }
        }
        
        try:
            with open(self.config_path, 'w') as f:
                json.dump(default_config, f, indent=2)
            logger.info(f"Created default MCP configuration with DeepWiki and Semgrep at {self.config_path}")
            # Update internal config
            self.servers_config = default_config
        except Exception as e:
            logger.error(f"Failed to create default config: {str(e)}")
            # Set default config in memory even if file creation fails
            self.servers_config = default_config
    
    def add_server(self, name: str, config: Dict[str, Any]) -> None:
        """
        Add a new MCP server configuration.
        
        Args:
            name: Server name
            config: Server configuration (command/args for stdio, url for http)
        """
        self.servers_config[name] = config
        self._save_config()
        logger.info(f"Added MCP server: {name}")
    
    def remove_server(self, name: str) -> bool:
        """
        Remove an MCP server configuration.
        
        Args:
            name: Server name to remove
            
        Returns:
            True if server was removed, False if not found
        """
        if name in self.servers_config:
            del self.servers_config[name]
            self._save_config()
            logger.info(f"Removed MCP server: {name}")
            return True
        return False
    
    def _save_config(self) -> None:
        """Save current configuration to file."""
        try:
            with open(self.config_path, 'w') as f:
                json.dump(self.servers_config, f, indent=2)
        except Exception as e:
            logger.error(f"Failed to save MCP configuration: {str(e)}")
    
    async def initialize(self) -> bool:
        """
        Initialize MCP client and connect to configured servers.
        
        Returns:
            True if initialization successful, False otherwise
        """
        if not MCP_AVAILABLE:
            logger.warning("MCP not available, skipping initialization")
            return False
        
        if not self.servers_config:
            logger.info("No MCP servers configured")
            return True
        
        # Filter out example/description entries
        active_servers = {
            name: config for name, config in self.servers_config.items()
            if not name.startswith('_')
        }
        
        if not active_servers:
            logger.info("No active MCP servers configured")
            return True
        
        try:
            self.client = MultiServerMCPClient(active_servers)
            self.tools = await self.client.get_tools()
            logger.info(f"Initialized MCP client with {len(self.tools)} tools from {len(active_servers)} servers")
            return True
        except Exception as e:
            logger.error(f"Failed to initialize MCP client: {str(e)}")
            return False
    
    async def cleanup(self) -> None:
        """Clean up MCP client connections."""
        if self.client:
            try:
                await self.client.close()
                logger.info("MCP client connections closed")
            except Exception as e:
                logger.error(f"Error closing MCP client: {str(e)}")
    
    def get_tools(self) -> List[Any]:
        """
        Get all available MCP tools.
        
        Returns:
            List of MCP tools
        """
        return self.tools
    
    def list_servers(self) -> Dict[str, Dict[str, Any]]:
        """
        List all configured MCP servers.
        
        Returns:
            Dictionary of server configurations
        """
        return {
            name: config for name, config in self.servers_config.items()
            if not name.startswith('_')
        }
    
    def get_server_info(self) -> Dict[str, Any]:
        """
        Get information about MCP integration status.
        
        Returns:
            Dictionary with integration status and server info
        """
        return {
            "mcp_available": MCP_AVAILABLE,
            "config_path": self.config_path,
            "servers_configured": len([k for k in self.servers_config.keys() if not k.startswith('_')]),
            "tools_available": len(self.tools),
            "client_initialized": self.client is not None
        }


# Global MCP integration instance
_mcp_integration: Optional[MCPIntegration] = None


def get_mcp_integration() -> MCPIntegration:
    """Get the global MCP integration instance."""
    global _mcp_integration
    if _mcp_integration is None:
        _mcp_integration = MCPIntegration()
    return _mcp_integration


async def initialize_mcp() -> bool:
    """Initialize MCP integration."""
    mcp = get_mcp_integration()
    return await mcp.initialize()


async def cleanup_mcp() -> None:
    """Cleanup MCP integration."""
    mcp = get_mcp_integration()
    await mcp.cleanup()


def get_mcp_tools() -> List[Any]:
    """Get all available MCP tools."""
    mcp = get_mcp_integration()
    return mcp.get_tools()


# Tools for agent use
def invoke():
    """Get this tool for LangGraph integration."""
    return mcp_server_tool


def mcp_server_tool(action: str, **kwargs) -> Dict[str, Any]:
    """
    MCP Server Management Tool for SWE Agent.
    
    This tool allows agents to manage MCP server connections and get information
    about available external tools.
    
    Args:
        action: Action to perform - 'list', 'info', 'add', 'remove', 'tools'
        **kwargs: Additional arguments based on action
        
    Returns:
        Dictionary with operation result
        
    Example usage:
        mcp_server_tool("list")  # List configured servers
        mcp_server_tool("info")  # Get integration status
        mcp_server_tool("tools") # List available MCP tools
        mcp_server_tool("add", name="math", config={...})  # Add server
        mcp_server_tool("remove", name="math")  # Remove server
    """
    try:
        mcp = get_mcp_integration()
        
        if action == "list":
            servers = mcp.list_servers()
            return {
                "success": True,
                "servers": servers,
                "count": len(servers)
            }
        
        elif action == "info":
            info = mcp.get_server_info()
            return {
                "success": True,
                "info": info
            }
        
        elif action == "tools":
            tools = mcp.get_tools()
            tool_info = []
            for tool in tools:
                tool_info.append({
                    "name": getattr(tool, 'name', 'Unknown'),
                    "description": getattr(tool, 'description', 'No description'),
                    "parameters": getattr(tool, 'parameters', {})
                })
            return {
                "success": True,
                "tools": tool_info,
                "count": len(tools)
            }
        
        elif action == "add":
            name = kwargs.get('name')
            config = kwargs.get('config')
            if not name or not config:
                return {
                    "success": False,
                    "error": "Both 'name' and 'config' are required for add action"
                }
            mcp.add_server(name, config)
            return {
                "success": True,
                "message": f"Added MCP server: {name}"
            }
        
        elif action == "remove":
            name = kwargs.get('name')
            if not name:
                return {
                    "success": False,
                    "error": "'name' is required for remove action"
                }
            removed = mcp.remove_server(name)
            return {
                "success": removed,
                "message": f"Removed MCP server: {name}" if removed else f"Server not found: {name}"
            }
        
        else:
            return {
                "success": False,
                "error": f"Unknown action: {action}. Available: list, info, tools, add, remove"
            }
    
    except Exception as e:
        logger.error(f"MCP server tool error: {str(e)}")
        return {
            "success": False,
            "error": f"MCP operation failed: {str(e)}"
        }


# For backward compatibility
mcp_server_management = mcp_server_tool