import json
import os
from typing import Literal
from langchain.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
from langchain.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.graph import END, START, MessagesState, StateGraph

CALLS: list[dict[str, object]] = []

@tool
def lookup_stock(sku: str) -> dict[str, object]:
    """Read the fixed demo stock for one SKU. This tool cannot change inventory."""
    result: dict[str, object] = {"sku": sku, "stock": {"DEMO-RED": 7, "DEMO-BLUE": 0}.get(sku), "source": "local-fixture"}
    CALLS.append(result)
    return result

llm = ChatOpenAI(
    model=os.environ.get("QUICKROUTER_MODEL", "glm-5.3"),
    api_key=os.environ["QUICKROUTER_API_KEY"],
    base_url="https://api.quickrouter.ai/v1",
    use_responses_api=False,
    max_tokens=4096,
    max_retries=0,
    timeout=90,
).bind_tools([lookup_stock])

def model_node(state: MessagesState) -> dict[str, list[AIMessage]]:
    answer = llm.invoke([SystemMessage(content="必须先用 lookup_stock 查询两个 SKU，再根据工具结果回答库存。工具返回 null 时说明未知，不能编造。")]+state["messages"])
    print("usage:", answer.usage_metadata)
    return {"messages": [answer]}

def tool_node(state: MessagesState) -> dict[str, list[ToolMessage]]:
    last = state["messages"][-1]
    if not isinstance(last, AIMessage):
        raise RuntimeError("Expected AIMessage")
    replies: list[ToolMessage] = []
    for call in last.tool_calls:
        if call["name"] != "lookup_stock":
            raise RuntimeError("Unexpected tool name")
        result = lookup_stock.invoke(call["args"])
        replies.append(ToolMessage(content=json.dumps(result, ensure_ascii=False), tool_call_id=call["id"]))
    return {"messages": replies}

def route(state: MessagesState) -> Literal["tools", "__end__"]:
    last = state["messages"][-1]
    return "tools" if isinstance(last, AIMessage) and last.tool_calls else END

builder = StateGraph(MessagesState)
builder.add_node("model", model_node)
builder.add_node("tools", tool_node)
builder.add_edge(START, "model")
builder.add_conditional_edges("model", route, ["tools", END])
builder.add_edge("tools", "model")
graph = builder.compile()
result = graph.invoke({"messages": [HumanMessage(content="查询 DEMO-RED 和 DEMO-BLUE 的库存，分别告诉我能否购买。只查库存。") ]}, config={"recursion_limit": 8})
found = {str(row["sku"]): row["stock"] for row in CALLS}
if found.get("DEMO-RED") != 7 or found.get("DEMO-BLUE") != 0:
    raise RuntimeError("Missing tool calls: a plausible text answer is not a pass")
print("tool_calls:", json.dumps(CALLS, ensure_ascii=False))
print("answer:", result["messages"][-1].content)