"""
第 7 步：完整作业版。

目标：
用 LangGraph 构建一个多角色协作的旅行规划 Agent，
并通过 MCP 动态调用本地工具，生成包含交通、住宿、景点、预算的旅行方案。

运行：
python step7_langgraph_travel_homework.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 的状态。

    你可以把它理解成一张“协作记录表”：
    每个角色完成自己的部分后，都把结果写进这张表。
    """

    user_request: str
    transport_plan: str
    hotel_plan: str
    attraction_plan: str
    budget_plan: str
    final_answer: str


def build_qwen_model() -> ChatOpenAI:
    """
    创建 Qwen 模型。

    这里使用的是 DashScope 的 OpenAI 兼容接口，
    所以可以用 LangChain 的 ChatOpenAI 来调用 Qwen。
    """
    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():
    """
    连接本地 MCP Server，并动态发现工具。

    这一步体现 MCP 的核心能力：
    Agent 主程序没有手写 search_hotels、search_transport 等函数，
    而是通过 MCP Client 从 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 工具，
    但 system_prompt 会限制它只做自己的专业工作。
    """
    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("完整作业版可用 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 attraction_node(state: TravelState) -> TravelState:
        answer = await call_role(
            model=model,
            tools=tools,
            system_prompt=(
                "你是景点体验规划师。"
                "你只负责景点、美食、夜景和每日体验安排。"
                "必须优先调用 recommend_attractions 工具。"
                "注意不要把行程排得太满。"
            ),
            user_request=state["user_request"],
        )
        return {"attraction_plan": answer}

    async def budget_node(state: TravelState) -> TravelState:
        answer = await call_role(
            model=model,
            tools=tools,
            system_prompt=(
                "你是预算规划师。"
                "你只负责估算总预算，并拆分交通、住宿、餐饮、市内交通、门票等项目。"
                "必须优先调用 estimate_budget 工具。"
            ),
            user_request=state["user_request"],
        )
        return {"budget_plan": answer}

    async def final_node(state: TravelState) -> TravelState:
        prompt = f"""
你是总规划师。请把四个角色的结果整合成一份可交作业的旅行规划方案。

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

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

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

景点体验规划师结果：
{state["attraction_plan"]}

预算规划师结果：
{state["budget_plan"]}

请按下面格式输出：

一、需求理解
二、交通方案
三、住宿方案
四、每日行程安排
五、预算估算
六、注意事项
七、技术说明：本项目如何体现 MCP 和 LangGraph 多角色协作

技术说明部分必须明确写出：
- MCP Server 暴露了哪些工具
- MCP Client 如何动态发现工具
- Function Calling 和 MCP 分别负责什么
- LangGraph 中每个角色节点负责什么
"""
        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("attraction_planner", attraction_node)
    graph.add_node("budget_planner", budget_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", "attraction_planner")
    graph.add_edge("attraction_planner", "budget_planner")
    graph.add_edge("budget_planner", "final_planner")
    graph.add_edge("final_planner", END)

    app = graph.compile()

    result = await app.ainvoke(
        {
            "user_request": (
                "我想 2026-08-15 从成都去重庆玩 3 天 2 晚，"
                "1 个人，预算控制在 1500 元以内。"
                "我喜欢美食、夜景和轻松一点的城市漫步，"
                "请生成交通、住宿、景点和预算方案。"
            )
        }
    )

    print("\n========== 完整作业最终结果 ==========\n")
    print(result["final_answer"])


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