import asyncio
from langchain_mcp_adapters.client import MultiServerMCPClient
from langchain_core.tools import StructuredTool


def wrap_mcp_tool(mcp_tool):
    """Wrapper to adapt MCP tool to LangChain format"""

    async def wrapped_func(**kwargs):
        # Just pass everything through, including tool_call_id
        result = await mcp_tool.ainvoke(kwargs)
        return result

    # Create a new StructuredTool with the wrapper
    return StructuredTool.from_function(
        coroutine=wrapped_func,
        name=mcp_tool.name,
        description=mcp_tool.description,
        args_schema=mcp_tool.args_schema if hasattr(mcp_tool, 'args_schema') else None
    )


async def load_boss_tools():
    return await boss_client.get_tools()


async def load_captain_tools():
    return await captain_client.get_tools()


async def load_tools():
    boss_tools = await load_boss_tools()
    captain_tools = await load_captain_tools()

    # wrap_mcp_tool is not async, so no await needed
    boss_agent_tools = [wrap_mcp_tool(tool) for tool in boss_tools]
    captain_agent_tools = [wrap_mcp_tool(tool) for tool in captain_tools]

    return boss_agent_tools, captain_agent_tools


boss_client = MultiServerMCPClient(
    {
        "BossServer": {
            "command": "python",
            "args": ["-m", "mcp_servers.boss_server"],
            "transport": "stdio",
        }
    }
)

captain_client = MultiServerMCPClient(
    {
        "CaptainServer": {
            "command": "python",
            "args": ["-m", "mcp_servers.captain_server"],
            "transport": "stdio",
        }
    }
)

boss_agent_tools, captain_agent_tools = asyncio.run(load_tools())