import os
from smolagents.mcp_client import MCPClient
from mcpadapt.smolagents_adapter import SmolAgentsAdapter
from mcpadapt.core import MCPAdapt
from mcp import StdioServerParameters
from mcp.types import Tool
from typing import Callable, Any
import smolagents
from functools import partial

# Custom adapter to fix tool names
class CustomSmolAgentsAdapter(SmolAgentsAdapter):
    def adapt(
        self, 
        func: Callable[[dict | None], Any], 
        mcp_tool: Tool
    ) -> smolagents.Tool:
        # Fix the tool name before passing it to the original adapter
        fixed_tool = Tool(
            name=mcp_tool.name.replace('-', '_'),
            description=mcp_tool.description,
            inputSchema=mcp_tool.inputSchema
        )
        return super().adapt(func, fixed_tool)

# Custom client that uses our adapter
def create_mcp_client(server_params):
    adapter = MCPAdapt(server_params, CustomSmolAgentsAdapter())
    
    # Create a client-like object with the same interface
    class CustomMCPClient:
        def __init__(self):
            self._tools = None
            self.connect()
            
        def connect(self):
            self._tools = adapter.__enter__()
            
        def disconnect(self, *args, **kwargs):
            adapter.__exit__(None, None, None)
            
        def get_tools(self):
            return self._tools
    
    return CustomMCPClient()