import asyncio
from typing import Any, List, Dict
from contextlib import AsyncExitStack
from pydantic import BaseModel, ConfigDict
from agenticblocks.core.block import Block

try:
    from mcp import ClientSession, StdioServerParameters
    from mcp.client.stdio import stdio_client
except ImportError:
    pass

class MCPProxyInput(BaseModel):
    model_config = ConfigDict(extra='allow')
    # Captures any arguments generated by the LLM

class MCPProxyOutput(BaseModel):
    result: Any

class MCPProxyBlock(Block[MCPProxyInput, MCPProxyOutput]):
    is_mcp_proxy: bool = True
    raw_mcp_schema: dict = {}
    session: Any = None
    
    async def run(self, input: MCPProxyInput) -> MCPProxyOutput:
        # Use normal dict() or model_dump(exclude_unset=True) to extract extras.
        # Note: model_dump may not spread "extras" to the root in some Pydantic versions,
        # but __dict__ with `extra='allow'` resolves this.
        args = {k: getattr(input, k) for k in input.model_fields_set | set(input.model_extra or {})}
        
        # Issues the Remote Procedure Call to the real MCP Server.
        result = await self.session.call_tool(self.name, args)
        
        # In MCP the response contains .content (Array of TextContent/ImageContent, etc.)
        # Converting the content remotely returned by the Server.
        text_responses = []
        if result and hasattr(result, "content"):
            for c in result.content:
                if getattr(c, "type", "") == "text":
                    text_responses.append(c.text)
        
        return MCPProxyOutput(result=text_responses)

class MCPClientBridge:
    """
    Transparent abstraction for fetching tools from compatible MCP servers.
    Transforms network tools into "ProxyBlocks" that can be plugged into Agents.
    """
    def __init__(self, command: str, args: List[str], env: Dict[str, str] = None):
        self.command = command
        self.args = args
        self.env = env
        self.exit_stack = AsyncExitStack()
        self.session = None

    async def connect(self) -> List[MCPProxyBlock]:
        """Establishes an iterative stdio link with the MCP Server and extracts its Tools."""
        server_params = StdioServerParameters(command=self.command, args=self.args, env=self.env)
        try:
            stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
        except Exception as e:
            raise RuntimeError(f"Failed to start MCP subprocess. Make sure '{self.command}' is installed. Error: {e}")
            
        self.read, self.write = stdio_transport
        
        self.session = await self.exit_stack.enter_async_context(ClientSession(self.read, self.write))
        await self.session.initialize()
        
        # Fetch tools from the server remotely
        response = await self.session.list_tools()
        
        proxy_blocks = []
        for tool in response.tools:
            # For each native tool from the MCP server, instantiate a proxy
            block = MCPProxyBlock(
                name=tool.name,
                description=tool.description or ""
            )
            block.raw_mcp_schema = tool.inputSchema
            block.session = self.session
            proxy_blocks.append(block)
            
        return proxy_blocks
        
    async def disconnect(self):
        """Closes the session safely."""
        await self.exit_stack.aclose()
