import asyncio
import json
import logging
from langchain.agents import create_tool_calling_agent, AgentExecutor
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from python_a2a.mcp import MCPClient
from langchain_core.tools import tool
from agent_learn.config import Config

conf = Config()

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 创建模型
llm = ChatOpenAI(base_url=conf.base_url,
                 api_key=conf.api_key,
                 model=conf.model_name,
                 temperature=0.1)

async def test_mcp_tools():
    # 连接到服务端，端口 8010
    client = MCPClient("http://localhost:8010")
    try:
        # 步骤 1：获取可用工具列表
        tools = await client.get_tools()
        logger.info("可用工具列表：")
        for tool_info in tools:
            print(tool_info)
            logger.info(f"- {tool_info.get('name', '未知')}: {tool_info.get('description', '无描述')}")

        # 用两个无参数函数，把 MCP 工具交给 Agent 使用
        @tool("get_weather")
        async def get_weather_tool() -> str:
            """获取教学示例中的北京天气，不查询实时天气。"""
            return await client.call_tool("get_weather")

        @tool("query_high_frequency_question")
        async def query_high_frequency_question_tool() -> str:
            """获取教学示例中的高频问题和答案。"""
            return await client.call_tool("query_high_frequency_question")

        tools = [get_weather_tool, query_high_frequency_question_tool]

        # 创建prompt模板
        prompt_template = ChatPromptTemplate.from_messages([
            ("system", "你是一个乐于助人的助手，能够调用工具回答用户问题。工具不需要传参。"),
            ("human", "{input}"),
            ("placeholder", "{agent_scratchpad}"),
        ])

        # 构建工具调用代理
        agent = create_tool_calling_agent(llm, tools, prompt_template)

        # 创建代理执行器
        agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

        # 代理调用
        print("MCP客户端启动，输入'quit'退出")
        while True:
            # 接收用户查询
            query = input("\nQuery: ").strip()
            if query.lower() == "quit":
                break
            # 发送用户查询给代理，并打印
            try:
                response = await agent_executor.ainvoke({"input": query})
                print(f"response-->{response}")
            except Exception as e:
                print(f"工具调用失败：{e}")
    except Exception as e:
        logger.error(f"MCP 客户端出错：{str(e)}", exc_info=True)
    finally:
        await client.close()

async def main():
    await test_mcp_tools()

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