import asyncio
import json

from langchain.agents import create_tool_calling_agent, AgentExecutor
from langchain_core.prompts import ChatPromptTemplate
from langchain_mcp_adapters.tools import load_mcp_tools
from langchain_openai import ChatOpenAI
from mcp import ClientSession
from mcp.client.streamable_http import streamablehttp_client

from conf import settings
from create_logger import logger

# 初始化LLM
llm = ChatOpenAI(
    model=settings.model_name,
    base_url=settings.base_url,
    api_key=settings.api_key,
    temperature=0.1
)


async def test_ticket_mcp():
    try:
        # 启动 MCP server，通过streamable建立连接
        async with streamablehttp_client("http://127.0.0.1:8001/mcp") as (read, write, _):
            # 使用读写通道创建 MCP 会话
            async with ClientSession(read, write) as session:
                try:
                    await session.initialize()
                    print("会话初始化成功，可以开始调用工具。")

                    # 从 session 自动获取 MCP server 提供的工具列表。
                    tools = await load_mcp_tools(session)
                    print(f"tools-->{tools}")

                    # 调用远程工具
                    # 测试1: 查询机票
                    sql_flights = "SELECT * FROM flight_tickets WHERE departure_city = '上海' AND arrival_city = '北京' AND DATE(departure_time) = '2025-10-28' AND cabin_type = '公务舱'"
                    result_flights = await session.call_tool("query_tickets", {"sql": sql_flights})
                    result_flights_data = json.loads(result_flights) if isinstance(result_flights, str) else result_flights
                    print(f"机票查询结果：{result_flights_data}")

                    # 测试2: 查询火车票
                    sql_trains = "SELECT * FROM train_tickets WHERE departure_city = '北京' AND arrival_city = '上海' AND DATE(departure_time) = '2025-10-22' AND seat_type = '二等座'"
                    result_trains = await session.call_tool("query_tickets", {"sql": sql_trains})
                    result_trains_data = json.loads(result_trains) if isinstance(result_trains, str) else result_trains
                    print(f"火车票查询结果：{result_trains_data}")

                    # 测试3: 查询演唱会票
                    sql_concerts = "SELECT * FROM concert_tickets WHERE city = '北京' AND artist = '刀郎' AND DATE(start_time) = '2025-10-31' AND ticket_type = '看台'"
                    result_concerts = await session.call_tool("query_tickets", {"sql": sql_concerts})
                    result_concerts_data = json.loads(result_concerts) if isinstance(result_concerts, str) else result_concerts
                    print(f"演唱会票查询结果：{result_concerts_data}")
                except Exception as e:
                    print(f"票务 MCP 测试出错：{str(e)}")
    except Exception as e:
        print(f"连接或会话初始化时发生错误: {e}")
        print("请确认服务端脚本已启动并运行在 http://127.0.0.1:8001/mcp")


async def test_weather_mcp():
    try:
        # 启动 MCP server，通过streamable建立连接
        async with streamablehttp_client("http://127.0.0.1:8002/mcp") as (read, write, _):
            # 使用读写通道创建 MCP 会话
            async with ClientSession(read, write) as session:
                try:
                    await session.initialize()
                    print("会话初始化成功，可以开始调用工具。")

                    # 从 session 自动获取 MCP server 提供的工具列表。
                    tools = await load_mcp_tools(session)
                    print(f"tools-->{tools}")

                    # 测试1: 查询指定日期天气
                    sql = "SELECT * FROM weather_data WHERE city = '北京' AND fx_date = '2025-10-28'"
                    result = await session.call_tool("query_weather", {"sql": sql})
                    result_data = json.loads(result) if isinstance(result, str) else result
                    print(f"指定日期天气结果：{result_data}")

                    # 测试2: 查询未来3天天气
                    sql_range = "SELECT * FROM weather_data WHERE city = '北京' AND fx_date BETWEEN '2025-10-28' AND '2025-10-30'"
                    result_range = await session.call_tool("query_weather", {"sql": sql_range})
                    result_range_data = json.loads(result_range) if isinstance(result_range, str) else result_range
                    print(f"天气范围查询结果：{result_range_data}")
                except Exception as e:
                    print(f"天气 MCP 测试出错：{str(e)}")
    except Exception as e:
        print(f"连接或会话初始化时发生错误: {e}")
        print("请确认服务端脚本已启动并运行在 http://127.0.0.1:8002/mcp")


async def order_tickets(query):
    try:
        # 启动 MCP server，通过streamable建立连接
        async with streamablehttp_client("http://127.0.0.1:8003/mcp") as (read, write, _):
            # 使用读写通道创建 MCP 会话
            async with ClientSession(read, write) as session:
                try:
                    await session.initialize()

                    # 从 session 自动获取 MCP server 提供的工具列表。
                    tools = await load_mcp_tools(session)
                    # print(f"tools-->{tools}")

                    # 创建 agent 的提示模板
                    prompt = ChatPromptTemplate.from_messages([
                        ("system",
                         "你是一个票务预定助手，能够调用工具来完成火车票、飞机票或演出票的预定。你需要仔细分析工具需要的参数，然后从用户提供的信息中提取信息。如果用户提供的信息不足以提取到调用工具所有必要参数，则向用户追问，以获取该信息。不能自己编撰参数。"),
                        ("human", "{input}"),
                        ("placeholder", "{agent_scratchpad}"),
                    ])

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

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

                    # 代理调用
                    response = await agent_executor.ainvoke({"input": query})

                    return response['output']
                except Exception as e:
                    logger.info(f"票务 MCP 测试出错：{str(e)}")
                    return f"票务 MCP 查询出错：{str(e)}"
    except Exception as e:
        logger.error(f"连接或会话初始化时发生错误: {e}")
        return "连接或会话初始化时发生错误"


if __name__ == "__main__":

    asyncio.run(test_ticket_mcp())

    asyncio.run(test_weather_mcp())

    while True:
        query = input("请输入查询：")
        if query == "exit":
            break
        print(asyncio.run(order_tickets(query)))