import asyncio
from typing import List, Dict, Any, Optional
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from .tools import Tool

class MCPServerClient:
    """
    Connects to a standard Model Context Protocol (MCP) server
    and generates native apexagents.Tool instances dynamically.
    """
    
    def __init__(self, command: str, args: List[str], env: Optional[Dict[str, str]] = None):
        """
        Task: Initializes the MCP Server Client parameters to connect to a specific local or remote stdio-based MCP server.
        
        Input Parameters:
            command (str): The executable command to launch the server (e.g. 'npx', 'uvx', 'python').
            args (List[str]): Arguments passed to the command (e.g. ['mcp-server-fetch']).
            env (Optional[Dict[str, str]]): An optional dictionary of environment variables specific to the server.
            
        Output Parameters / Returns: None
        
        Raises: None
        """
        self.server_parameters = StdioServerParameters(
            command=command,
            args=args,
            env=env
        )
        self.session: Optional[ClientSession] = None
        self._exit_stack = None
        
    async def connect(self):
        """
        Task: Initializes the asynchronous stdio connection to the MCP server.
        
        Input Parameters: None
            
        Output Parameters / Returns: None
        
        Raises:
            Exception: If the underlying MCP connection fails to initialize.
        """
        from contextlib import AsyncExitStack
        
        self._exit_stack = AsyncExitStack()
        read, write = await self._exit_stack.enter_async_context(
            stdio_client(self.server_parameters)
        )
        self.session = await self._exit_stack.enter_async_context(
            ClientSession(read, write)
        )
        await self.session.initialize()
        
    async def close(self):
        """
        Task: Closes the connection to the MCP server safely, ensuring all background tasks are exited.
        
        Input Parameters: None
            
        Output Parameters / Returns: None
        
        Raises: None
        """
        if self._exit_stack:
            await self._exit_stack.aclose()
            
    async def get_tools(self) -> List[Tool]:
        """
        Task: Fetches all available tools from the connected MCP server and wraps them into native apexagents Tool instances for seamless LLM execution.
        
        Input Parameters: None
            
        Output Parameters / Returns:
            List[Tool]: An array of natively wrapped tools ready to be passed to an Agent.
            
        Raises:
            RuntimeError: If called prior to successfully awaiting `.connect()`.
        """
        if not self.session:
            raise RuntimeError("MCPServerClient is not connected. Call await client.connect() first.")
            
        mcp_tools_response = await self.session.list_tools()
        apex_tools = []
        
        for mcp_tool in mcp_tools_response.tools:
            # Create a dynamic function wrapper for executing the MCP tool natively
            def make_executor(tool_name: str):
                async def execute_mcp_tool(**kwargs) -> str:
                    result = await self.session.call_tool(tool_name, arguments=kwargs)
                    
                    # MCP returns a list of contents (text, images, etc.)
                    # Extract the text chunks to pass back to the LLM memory
                    texts = []
                    for content in result.content:
                        if getattr(content, "type", "") == "text":
                            texts.append(content.text)
                    return "\n".join(texts)
                
                # Override the name and docstring so Tool() parses it correctly
                execute_mcp_tool.__name__ = tool_name
                execute_mcp_tool.__doc__ = mcp_tool.description or "External MCP Tool."
                
                return execute_mcp_tool
                
            dynamic_func = make_executor(mcp_tool.name)
            
            # Create the apex tool
            new_tool = Tool(dynamic_func)
            
            # Since MCP tools dictate their own schema explicitly, we override our generated
            # schema with the exact authoritative schema from the MCP server
            new_tool.schema = {
                "type": "function",
                "function": {
                    "name": mcp_tool.name,
                    "description": mcp_tool.description or "",
                    "parameters": mcp_tool.inputSchema
                }
            }
            
            apex_tools.append(new_tool)
            
        return apex_tools
