# examples/mcp_modes/plan_and_execute.py
import os, asyncio
from typing import List
from langchain_core.messages import HumanMessage, ToolMessage
from langgraph_qwen.chat_model import ChatQwenOpenAICompat
from langgraph_qwen.mcp import load_mcp_tools

async def plan(task: str) -> List[str]:
    planner = ChatQwenOpenAICompat(temperature=0)
    ai = await planner.ainvoke([HumanMessage(content=f"把任务拆成可执行步骤（每行一步）：\n{task}")])
    steps = [s.strip() for s in str(ai.content).splitlines() if s.strip()]
    return steps[:8]

async def tools_for_step(step_idx: int):
    # ★ 示例：偶数步用 weather，奇数步用 play/test 工具；也可来自文件
    if step_idx % 2 == 0:
        return await load_mcp_tools(servers={
            "weather": {"url":"http://127.0.0.1:8000/mcp/","transport":"streamable_http"}
        })
    else:
        cfg = os.getenv("MCP_CONFIG_PATH")  # 比如 ./mcp_servers.yaml
        return await load_mcp_tools(config_path=cfg)

async def execute(steps: List[str], max_tool_steps_per_step: int = 4):
    msgs = []
    for i, step in enumerate(steps, 1):
        tools = await tools_for_step(i)
        tool_map = {t.name: t for t in tools}
        model = ChatQwenOpenAICompat(temperature=0).bind_tools(tools).bind(tool_choice="auto")

        msgs.append(HumanMessage(content=f"执行第{i}步：{step}"))
        for _ in range(max_tool_steps_per_step):
            ai = await model.ainvoke(msgs)
            msgs.append(ai)
            calls = getattr(ai, "tool_calls", []) or ai.additional_kwargs.get("tool_calls", [])
            if not calls:
                break
            for call in calls:
                name, args, call_id = call.get("name"), call.get("args", {}), call.get("id") or ""
                tool = tool_map.get(name)
                if not tool:
                    msgs.append(ToolMessage(tool_call_id=call_id, content=f"Unknown tool: {name}"))
                    continue
                try:
                    out = await tool.ainvoke(args) if hasattr(tool, "ainvoke") else tool.invoke(args)
                    msgs.append(ToolMessage(tool_call_id=call_id, content=str(out)))
                except Exception as e:
                    msgs.append(ToolMessage(tool_call_id=call_id, content=f"Error: {e}"))

    final = await ChatQwenOpenAICompat(temperature=0).ainvoke(
        msgs + [HumanMessage(content="请汇总执行结果，简洁给出结论。")]
    )
    print("=== Final ===")
    print(final.content)

async def main():
    steps = await plan("在巴塞罗那找一个带游泳池的民宿，然后搜索附近的餐厅和景点")
    await execute(steps, max_tool_steps_per_step=4)

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