"""LangGraph 工具调用示例：天气查询 + 计算器。

图结构：

    START -> model
               | 有 tool_calls
               v
             tools
               |
               +---------> model

             model -- 无 tool_calls --> END

天气数据是本地模拟数据，重点是学习 LangGraph 的工具调用循环。
"""

import os
from typing import Literal

from dotenv import load_dotenv
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.graph import END, START, MessagesState, StateGraph
from langgraph.prebuilt import ToolNode

load_dotenv()


# 1. 定义工具。类型注解和 docstring 会变成提供给模型的工具 schema。
@tool
def get_weather(city: str) -> str:
    """查询指定城市当前的天气。城市名称使用中文，例如北京、上海。"""
    weather_data = {
        "北京": "晴，25°C，北风 2 级",
        "上海": "多云，28°C，东南风 3 级",
        "广州": "雷阵雨，30°C，南风 2 级",
        "深圳": "阴，29°C，东风 2 级",
    }
    return weather_data.get(city, f"暂时没有 {city} 的天气数据")


@tool
def calculator(a: float, b: float, operation: Literal["add", "subtract", "multiply", "divide"]) -> float:
    """对两个数字进行加、减、乘、除运算。

    Args:
        a: 第一个数字。
        b: 第二个数字。
        operation: 运算类型：add、subtract、multiply 或 divide。
    """
    if operation == "add":
        return a + b
    if operation == "subtract":
        return a - b
    if operation == "multiply":
        return a * b
    if b == 0:
        raise ValueError("除数不能为 0")
    return a / b


# 2. 将工具绑定给模型。绑定后，模型可以在 AIMessage.tool_calls 中请求工具。
tools = [get_weather, calculator]
model = ChatOpenAI(
    model=os.getenv("OPENAI_MODEL", "openai/gpt-5.4"),
    api_key=os.getenv("OPENAI_API_KEY"),
    base_url=os.getenv("OPENAI_BASE_URL", "https://api.qnaigc.com/v1"),
    temperature=0,
    max_retries=2,
)
model_with_tools = model.bind_tools(tools)


# 3. 模型节点：读取完整消息历史，返回一条新的 AIMessage。
def call_model(state: MessagesState) -> dict:
    response = model_with_tools.invoke(state["messages"])
    return {"messages": [response]}


# 4. 路由函数：检查模型是否发起了工具调用。
def route_after_model(state: MessagesState) -> Literal["tools", "__end__"]:
    last_message = state["messages"][-1]
    if last_message.tool_calls:
        return "tools"
    return END


# 5. ToolNode 自动完成：找到工具、传入参数、执行工具、生成 ToolMessage。
tool_node = ToolNode(tools)

builder = StateGraph(MessagesState)
builder.add_node("model", call_model)
builder.add_node("tools", tool_node)

builder.add_edge(START, "model")
builder.add_conditional_edges(
    "model",
    route_after_model,
    {
        "tools": "tools",
        END: END,
    },
)
builder.add_edge("tools", "model")

graph = builder.compile()


def ask(question: str) -> str:
    """调用助手，并打印每个节点产生的更新。"""
    final_answer = ""

    for event in graph.stream(
        {"messages": [{"role": "user", "content": question}]},
        stream_mode="updates",
    ):
        for node_name, update in event.items():
            print(f"\n[{node_name}]")
            for message in update.get("messages", []):
                message.pretty_print()
                if node_name == "model" and not message.tool_calls:
                    final_answer = str(message.content)

    return final_answer


def main() -> None:
    print("=== 天气 + 计算器工具助手 ===")
    print("输入 quit 退出。")

    while True:
        question = input("\n用户：").strip()
        if question.lower() in {"quit", "exit", "q"}:
            break
        if not question:
            continue

        answer = ask(question)
        print(f"\n最终回答：{answer}")


if __name__ == "__main__":
    main()