"""
第 6 步：用 LangGraph 做一个最小多角色旅行规划 Agent。

这一版先只做两个角色：
1. 交通规划师：负责交通
2. 住宿规划师：负责住宿
3. 总结节点：把两个角色结果合成一段旅行建议

运行：
python step6_langgraph_two_roles.py
"""

from __future__ import annotations

import asyncio
import os
import sys
from pathlib import Path
from typing import TypedDict

from dotenv import load_dotenv
from langchain.agents import create_agent
from langchain_mcp_adapters.client import MultiServerMCPClient
from langchain_openai import ChatOpenAI
from langgraph.graph import END, START, StateGraph


class TravelState(TypedDict, total=False):
    """
    LangGraph 里的 state，可以理解成“多个角色之间传递的工作记录”。
    """

    user_request: str
    transport_plan: str
    hotel_plan: str
    final_answer: str


def build_qwen_model() -> ChatOpenAI:
    load_dotenv()

    api_key = os.getenv("DASHSCOPE_API_KEY")
    if not api_key:
        raise RuntimeError("请先在 .env 里填写 DASHSCOPE_API_KEY。")

    return ChatOpenAI(
        api_key=api_key,
        base_url=os.getenv("QWEN_BASE_URL", "https://dashscope.aliyuncs.com/compatible-mode/v1"),
        model=os.getenv("QWEN_MODEL", "qwen-plus"),
        temperature=0.2,
    )


async def build_mcp_tools():
    """
    连接 my_mcp_demo.py 这个本地 MCP Server，并动态获取工具列表。
    """
    server_file = Path(__file__).with_name("my_mcp_demo.py").resolve()

    client = MultiServerMCPClient(
        {
            "travel-tools": {
                "transport": "stdio",
                "command": sys.executable,
                "args": [str(server_file)],
            }
        }
    )

    return await client.get_tools()


async def call_role(model: ChatOpenAI, tools, system_prompt: str, user_request: str) -> str:
    """
    创建一个角色 Agent，让它根据用户需求调用 MCP 工具。
    """
    agent = create_agent(
        model=model,
        tools=tools,
        system_prompt=system_prompt,
    )

    result = await agent.ainvoke(
        {
            "messages": [
                {
                    "role": "user",
                    "content": user_request,
                }
            ]
        }
    )

    return result["messages"][-1].content


async def main() -> None:
    model = build_qwen_model()
    tools = await build_mcp_tools()

    print("LangGraph 可用 MCP 工具：", [tool.name for tool in tools])

    async def transport_node(state: TravelState) -> TravelState:
        """
        第一个角色：交通规划师。
        """
        answer = await call_role(
            model=model,
            tools=tools,
            system_prompt=(
                "你是交通规划师。"
                "你只负责解决城市间交通问题。"
                "如果用户提到从哪里到哪里，请优先调用 search_transport 工具。"
                "回答要简洁，说明推荐交通方式和理由。"
            ),
            user_request=state["user_request"],
        )
        return {"transport_plan": answer}

    async def hotel_node(state: TravelState) -> TravelState:
        """
        第二个角色：住宿规划师。
        """
        answer = await call_role(
            model=model,
            tools=tools,
            system_prompt=(
                "你是住宿规划师。"
                "你只负责解决住哪里、住几晚、预算多少的问题。"
                "如果用户提到住宿或几天几晚，请优先调用 search_hotels 工具。"
                "回答要简洁，说明住宿预算和区域建议。"
            ),
            user_request=state["user_request"],
        )
        return {"hotel_plan": answer}

    async def final_node(state: TravelState) -> TravelState:
        """
        第三个节点：总规划师。
        它不直接调 MCP 工具，只负责汇总前面两个角色的结果。
        """
        prompt = f"""
请把下面两个角色的结果整合成一份简洁旅行建议。

用户需求：
{state["user_request"]}

交通规划师结果：
{state["transport_plan"]}

住宿规划师结果：
{state["hotel_plan"]}

请输出：
1. 交通建议
2. 住宿建议
3. 一句话总结
"""
        response = await model.ainvoke(prompt)
        return {"final_answer": response.content}

    graph = StateGraph(TravelState)

    graph.add_node("transport_planner", transport_node)
    graph.add_node("hotel_planner", hotel_node)
    graph.add_node("final_planner", final_node)

    graph.add_edge(START, "transport_planner")
    graph.add_edge("transport_planner", "hotel_planner")
    graph.add_edge("hotel_planner", "final_planner")
    graph.add_edge("final_planner", END)

    app = graph.compile()

    result = await app.ainvoke(
        {
            "user_request": (
                "我想 2026-08-15 从成都去重庆玩 3 天 2 晚，"
                "预算控制在 1500 元以内，帮我安排交通和住宿。"
            )
        }
    )

    print("\n========== LangGraph 最终结果 ==========\n")
    print(result["final_answer"])


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