import asyncio
import logging

# from langchain.agents import create_react_agent
from langgraph.prebuilt import create_react_agent
from langchain_openai import ChatOpenAI
from langchain_mcp_adapters.client import MultiServerMCPClient
import os
import dotenv

logging.basicConfig(
    level=logging.DEBUG,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)

dotenv.load_dotenv()

## 设置环境变量
os.environ['OPENAI_API_KEY'] = os.getenv("OLLAMA_API_KEY")
os.environ['OPENAI_BASE_URL'] = os.getenv("OLLAMA_BASE_URL")

# 默认的 'model_name': 'deepseek-ai/DeepSeek-V3.1',
llm = ChatOpenAI(model="gemma4:26b")


def print_optimized_result(agent_response):
    """
        解析代理响应并输出优化后的结果。
        :param agent_response: 代理返回的完整响应
        """
    messages = agent_response.get("messages", [])
    steps = []  # 用于记录计算步骤
    final_answer = None  # 最终答案

    for message in messages:
        if hasattr(message, "additional_kwargs") and "tool_calls" in message.additional_kwargs:
            # 提取工具调用信息
            tool_calls = message.additional_kwargs["tool_calls"]
            for tool_call in tool_calls:
                tool_name = tool_call["function"]["name"]
                tool_args = tool_call["function"]["arguments"]
                steps.append(f"调用工具: {tool_name}({tool_args})")
        elif message.type == "tool":
            # 提取工具执行结果
            tool_name = message.name
            tool_result = message.content
            steps.append(f"{tool_name} 的结果是: {tool_result}")
        elif message.type == "ai":
            # 提取最终答案
            final_answer = message.content

    # 打印优化后的结果
    print("\n计算过程:")
    for step in steps:
        print(f"- {step}")
    if final_answer:
        print(f"\n最终答案: {final_answer}")


async def execute():
    # 1. 创建langchain中的mcp客户端  ——  uv add langchain_mcp_adapters
    client = MultiServerMCPClient(
        #         mcp.run(transport="streamable-http", host="127.0.0.1", port=9000)
        {
            # 这里是定义服务端信息的，可以有多个服务端
            "weather": {
                "url": "http://localhost:9000/mcp",
                "transport": "streamable_http",
            },
            "math": {
                "command": "python",  # npx uvx
                "args": ["./mcp/math_server.py"],
                "transport": "stdio"
            }
        }
    )

    try:
        # 2. 通过客户端获取工具列表
        ## 这里是通过服务端获取，所以可能会有异常（比如服务端没有启动，或者网络连接有问题
        tools = await client.get_tools()

        # 3. 创建一个智能代理，能够完成 思考--> 行动 --> 观察 --> 思考 --> 行动 --> ... --> 最终答案
        agent = create_react_agent(model=llm, tools=tools)
        while True:
            user_input = input("请输入你的问题（输入exit则退出） > ")
            if user_input == "exit":
                print("感谢使用，再见👋🏻")
                break

            agent_response = await agent.ainvoke({"messages": user_input})
            ## agent会自己去调用 tools ，不需要我们去进行调用
            print_optimized_result(agent_response)

    finally:
        # 资源回收
        if hasattr(client, 'close'):
            client.close()


if __name__ == '__main__':
    asyncio.run(execute())