"""
MCP (Model Context Protocol) client for LangCode Agent
Enables integration with MCP servers like skills.sh to discover and use tools
"""
import json
import requests
from typing import Dict, List, Any, Optional
from langchain_core.tools import Tool
import logging

logger = logging.getLogger(__name__)

class MCPClient:
    """Client for communicating with MCP (Model Context Protocol) servers"""
    
    def __init__(self, server_url: str, headers: Optional[Dict[str, str]] = None, timeout: int = 30):
        """
        Initialize MCP client
        
        Args:
            server_url: URL of the MCP server
            headers: Optional headers for authentication
            timeout: Request timeout in seconds
        """
        self.server_url = server_url.rstrip('/')
        self.headers = headers or {}
        self.timeout = timeout
        self._available_tools: Optional[List[Dict[str, Any]]] = None
    
    def _make_request(self, method: str, endpoint: str, data: Optional[Dict] = None) -> Dict[str, Any]:
        """
        Make a JSON-RPC request to the MCP server
        
        Args:
            method: HTTP method (GET, POST, etc.)
            endpoint: API endpoint
            data: Optional data to send
            
        Returns:
            Response data as dictionary
            
        Raises:
            Exception: If request fails
        """
        url = f"{self.server_url}{endpoint}"
        
        # For MCP, we typically use JSON-RPC over HTTP POST
        if method.upper() == "POST":
            # JSON-RPC 2.0 format
            payload = {
                "jsonrpc": "2.0",
                "method": endpoint.lstrip('/'),  # Remove leading slash for method name
                "params": data or {},
                "id": 1
            }
            response = requests.post(
                url, 
                json=payload, 
                headers={**self.headers, "Content-Type": "application/json"},
                timeout=self.timeout
            )
        else:
            response = requests.request(
                method, 
                url, 
                headers=self.headers,
                params=data,
                timeout=self.timeout
            )
        
        response.raise_for_status()
        
        # Handle JSON-RPC response
        if method.upper() == "POST":
            result = response.json()
            if "error" in result:
                raise Exception(f"MCP error: {result['error']}")
            return result.get("result", {})
        else:
            return response.json()
    
    def list_tools(self) -> List[Dict[str, Any]]:
        """
        List available tools from the MCP server
        
        Returns:
            List of tool specifications
        """
        if self._available_tools is not None:
            return self._available_tools
            
        try:
            # MCP standard method for listing tools
            result = self._make_request("POST", "/tools/list")
            self._available_tools = result.get("tools", [])
            logger.info(f"Discovered {len(self._available_tools)} tools from MCP server {self.server_url}")
            return self._available_tools
        except Exception as e:
            logger.error(f"Failed to list tools from MCP server {self.server_url}: {e}")
            self._available_tools = []
            return []
    
    def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
        """
        Call a specific tool on the MCP server
        
        Args:
            tool_name: Name of the tool to call
            arguments: Arguments to pass to the tool
            
        Returns:
            Tool execution result
        """
        try:
            result = self._make_request("POST", "/tools/call", {
                "name": tool_name,
                "arguments": arguments
            })
            return result
        except Exception as e:
            logger.error(f"Failed to call tool {tool_name} on MCP server {self.server_url}: {e}")
            raise Exception(f"MCP tool call failed: {e}")
    
    def refresh_tools(self):
        """Refresh the cached list of available tools"""
        self._available_tools = None
        self.list_tools()


def mcp_to_langchain_tool(mcp_client: MCPClient, tool_spec: Dict[str, Any]) -> Tool:
    """
    Convert an MCP tool specification to a LangChain Tool
    
    Args:
        mcp_client: MCP client instance
        tool_spec: Tool specification from MCP server
        
    Returns:
        LangChain Tool object
    """
    def tool_function(*args, **kwargs):
        # Handle both positional and keyword arguments
        if args and len(args) == 1 and isinstance(args[0], dict):
            # Single dict argument passed positionally
            tool_args = args[0]
        elif kwargs:
            # Keyword arguments
            tool_args = kwargs
        else:
            # No arguments or multiple positional args (shouldn't happen for MCP tools)
            tool_args = {}
        
        try:
            result = mcp_client.call_tool(tool_spec["name"], tool_args)
            # Convert result to string if it isn't already
            if isinstance(result, str):
                return result
            else:
                return json.dumps(result, indent=2)
        except Exception as e:
            return f"Error calling MCP tool {tool_spec['name']}: {str(e)}"
    
    # Extract description, default to empty string if not provided
    description = tool_spec.get("description", "")
    if not description:
        description = f"MCP tool: {tool_spec['name']}"
    
    return Tool(
        name=tool_spec["name"],
        description=description,
        func=tool_function
    )


def get_mcp_tools_from_client(mcp_client: MCPClient) -> List[Tool]:
    """
    Get all available tools from an MCP client as LangChain Tools
    
    Args:
        mcp_client: Configured MCP client
        
    Returns:
        List of LangChain Tool objects
    """
    tools = []
    try:
        tool_specs = mcp_client.list_tools()
        for tool_spec in tool_specs:
            langchain_tool = mcp_to_langchain_tool(mcp_client, tool_spec)
            tools.append(langchain_tool)
        logger.info(f"Converted {len(tools)} MCP tools to LangChain Tools")
    except Exception as e:
        logger.error(f"Failed to get tools from MCP client: {e}")
    
    return tools