import asyncio
import sys
from pathlib import Path

from langchain.agents import create_agent
from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel
from langchain_core.messages import AIMessage, ToolMessage
from langchain_mcp_adapters.client import MultiServerMCPClient


class ToolCallingFakeModel(FakeMessagesListChatModel):
    def bind_tools(self, tools, *, tool_choice=None, **kwargs):
        return self


async def main():
    server_path = Path(__file__).with_name("math_mcp_server.py")
    client = MultiServerMCPClient(
        {
            "math": {
                "transport": "stdio",
                "command": sys.executable,
                "args": [str(server_path)],
            }
        }
    )

    tools = await client.get_tools()
    print("MCP tools:", ", ".join(tool.name for tool in tools))

    add_tool = next(tool for tool in tools if tool.name == "add")
    model = ToolCallingFakeModel(
        responses=[
            AIMessage(
                content="",
                tool_calls=[
                    {
                        "name": add_tool.name,
                        "args": {"a": 7, "b": 5},
                        "id": "call_add",
                    }
                ],
            ),
            AIMessage(content="The MCP add tool returned 12."),
        ]
    )

    agent = create_agent(model=model, tools=tools)
    result = await agent.ainvoke(
        {
            "messages": [
                {
                    "role": "user",
                    "content": "Use the MCP add tool to add 7 and 5.",
                }
            ]
        }
    )

    for message in result["messages"]:
        if isinstance(message, ToolMessage):
            content = message.content
            if isinstance(content, list):
                text_blocks = [
                    block.get("text", "")
                    for block in content
                    if isinstance(block, dict) and block.get("type") == "text"
                ]
                content = " ".join(text_blocks)
            print(f"Tool result ({message.name}): {content}")

    print("Agent reply:", result["messages"][-1].content)


if __name__ == "__main__":
    asyncio.run(main())