from __future__ import annotations
"""直接调用自建旅行MCP服务器工具，验证返回结构是否符合预期。"""
import asyncio
from typing import Any
from mcp import ClientSession
from mcp.shared.memory import create_client_server_memory_streams
from rich.console import Console
from rich.table import Table
from mcp_servers.travel_search_server import mcp


console = Console()

def get_payload(result: Any) -> dict[str, Any]:
    payload = result.structuredContent

    if not isinstance(payload, dict):
        raise RuntimeError(
            "MCP工具没有返回结构化结果。"
        )

    return payload


async def main() -> None:
    console.rule(
        "[bold blue]TravelMind MCP Server测试"
    )

    # MCP SDK 1.x 内存直连模式：
    # 1. 创建双向内存流
    # 2. 服务端通过 asyncio.create_task 在后台运行
    # 3. 客户端通过 ClientSession 连接并初始化
    async with create_client_server_memory_streams() as (
        client_streams,
        server_streams,
    ):
        client_read, client_write = client_streams
        server_read, server_write = server_streams

        init_opts = mcp._mcp_server.create_initialization_options()

        async with ClientSession(client_read, client_write) as client:
            server_task = asyncio.create_task(
                mcp._mcp_server.run(
                    server_read,
                    server_write,
                    init_opts,
                )
            )
            try:
                await client.initialize()

                tools_result = await client.list_tools()

                tool_table = Table(title="已发现MCP工具")
                tool_table.add_column("工具名")
                tool_table.add_column("说明")

                for tool in tools_result.tools:
                    tool_table.add_row(
                        tool.name,
                        (tool.description or "")[:100],
                    )

                console.print(tool_table)

                console.rule("[bold]1. 查询往返去程")

                outbound_result = await client.call_tool(
                    "search_flights",
                    {
                        "departure_airports": [
                            "SHA",
                            "PVG",
                        ],
                        "arrival_airports": [
                            "CTU",
                            "TFU",
                        ],
                        "outbound_date": "2026-08-10",
                        "return_date": "2026-08-13",
                        "flight_type": "round_trip",
                        "adults": 2,
                        "currency": "CNY",
                    },
                )

                outbound_payload = get_payload(
                    outbound_result
                )

                outbound_flights = outbound_payload.get(
                    "flights",
                    [],
                )

                console.print(
                    {
                        "去程候选总数": (
                            outbound_payload.get(
                                "total_count"
                            )
                        ),
                        "实际结构化数量": len(
                            outbound_flights
                        ),
                    }
                )

                if not outbound_flights:
                    raise RuntimeError(
                        "没有查询到去程航班。"
                    )

                selected_outbound = next(
                    (
                        flight
                        for flight in outbound_flights
                        if flight.get("departure_token")
                    ),
                    None,
                )

                if selected_outbound is None:
                    raise RuntimeError(
                        "去程候选中没有 departure_token。"
                    )

                console.print(
                    {
                        "测试选中的去程": (
                            selected_outbound.get(
                                "flight_numbers"
                            )
                        ),
                        "出发时间": (
                            selected_outbound.get(
                                "departure_time"
                            )
                        ),
                        "到达时间": (
                            selected_outbound.get(
                                "arrival_time"
                            )
                        ),
                        "价格": selected_outbound.get(
                            "price"
                        ),
                    }
                )

                console.rule("[bold]2. 查询对应返程")

                return_result = await client.call_tool(
                    "search_return_flights",
                    {
                        "departure_token": (
                            selected_outbound[
                                "departure_token"
                            ]
                        ),
                        "departure_airports": [
                            "SHA",
                            "PVG",
                        ],
                        "arrival_airports": [
                            "CTU",
                            "TFU",
                        ],
                        "outbound_date": "2026-08-10",
                        "return_date": "2026-08-13",
                        "adults": 2,
                        "currency": "CNY",
                    },
                )

                return_payload = get_payload(
                    return_result
                )

                console.print(
                    {
                        "返程候选总数": (
                            return_payload.get(
                                "total_count"
                            )
                        ),
                        "实际结构化数量": len(
                            return_payload.get(
                                "flights",
                                [],
                            )
                        ),
                    }
                )

                console.rule("[bold]3. 查询酒店")

                hotel_result = await client.call_tool(
                    "search_hotels",
                    {
                        "query": "成都酒店",
                        "check_in_date": "2026-08-10",
                        "check_out_date": "2026-08-13",
                        "adults": 2,
                        "currency": "CNY",
                    },
                )

                hotel_payload = get_payload(
                    hotel_result
                )

                hotels = hotel_payload.get("hotels", [])

                console.print(
                    {
                        "酒店候选总数": (
                            hotel_payload.get(
                                "total_count"
                            )
                        ),
                        "实际结构化数量": len(hotels),
                    }
                )

                if hotels:
                    console.print(
                        {
                            "首家酒店": hotels[0].get(
                                "name"
                            ),
                            "评分": hotels[0].get(
                                "overall_rating"
                            ),
                            "每晚价格": hotels[0].get(
                                "price_per_night"
                            ),
                            "总价格": hotels[0].get(
                                "total_price"
                            ),
                        }
                    )

            finally:
                server_task.cancel()
                try:
                    await server_task
                except asyncio.CancelledError:
                    pass

    console.rule(
        "[bold green]MCP Server测试通过"
    )


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