# In a real system this would be a database query or microservice call.
# The keys are order IDs; each value holds the order status and total amount.
_FAKE_ORDER_DB = {
    "A123": {"status": "out for delivery", "amount": 49.99},
    "B456": {"status": "delivered", "amount": 19.50},
}

from langchain_core.tools import tool

@tool
def get_order_status(order_id: str) -> str:
    """
    Look up the status of a customer order by its order ID.

    Note: errors are returned as strings (not raised as exceptions) so the LLM
    can reason about the failure and ask the user to retry, rather than
    crashing the agent loop.
    """
    order = _FAKE_ORDER_DB.get(order_id)
    if order is None:
        return f"Error: no order found with ID '{order_id}'. Ask the user to double check the ID."
    return f"Order {order_id} status: {order['status']}."


@tool
def issue_refund(order_id: str, amount: float) -> str:
    """
    Issue a refund for a given order ID and dollar amount.

    Tool-level guardrail: the amount cap is enforced here in Python, not
    in a prompt. Never rely on the LLM alone to enforce financial constraints —
    always validate them in code.
    """
    order = _FAKE_ORDER_DB.get(order_id)
    if order is None:
        return f"Error: cannot refund unknown order '{order_id}'."
    if amount > order["amount"]:
        return f"Error: refund amount ${amount} exceeds the original order amount ${order['amount']}."
    return f"Refunded ${amount} for order {order_id}."

from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langgraph.checkpoint.memory import MemorySaver
from langgraph.prebuilt import ToolNode, tools_condition  # prebuilt ReAct helpers
from typing import Annotated
from typing_extensions import TypedDict

# How many tool calls the agent may make in one session before being forced
# to produce a final answer. Without a cap a confused agent can loop forever,
# burning API tokens and never returning to the user.
MAX_TOOL_CALLS = 6

tools = [get_order_status, issue_refund]
llm = get_llm()

# bind_tools() teaches the LLM which tools exist and what arguments they accept.
# The model may now emit tool_calls in its response instead of (or in addition to) text.
llm_with_tools = llm.bind_tools(tools)


class State(TypedDict):
    messages: Annotated[list, add_messages]
    tool_call_count: int  # running total; compared against MAX_TOOL_CALLS


def agent_node(state: State) -> dict:
    tool_call_count = state.get("tool_call_count", 0)

    if tool_call_count >= MAX_TOOL_CALLS:
        # Force a final answer: use plain llm (no tools bound) so it cannot
        # emit tool_calls and loop further.
        forced = llm.invoke(
            state["messages"]
            + [("system", "You have used the maximum number of tool calls. Answer now without calling any more tools.")]
        )
        # tool_call_count is NOT incremented — this is the terminal reply.
        return {"messages": [forced]}

    response = llm_with_tools.invoke(state["messages"])
    # Only count as a tool call round if the model actually requested a tool.
    increment = 1 if getattr(response, "tool_calls", None) else 0
    return {"messages": [response], "tool_call_count": tool_call_count + increment}


# --------------- Graph wiring ---------------
graph_builder = StateGraph(State)
graph_builder.add_node("agent", agent_node)

# ToolNode is a prebuilt node that:
#   1. Inspects the last AIMessage for .tool_calls.
#   2. Calls each listed function with the provided arguments.
#   3. Wraps each result in a ToolMessage and appends it to state["messages"].
graph_builder.add_node("tools", ToolNode(tools))

graph_builder.add_edge(START, "agent")

# tools_condition is a prebuilt router:
#   - If the last AIMessage has .tool_calls → routes to "tools".
#   - If it has no tool_calls (plain answer) → routes to END.
# This is the core of the ReAct (Reasoning + Acting) loop.
graph_builder.add_conditional_edges("agent", tools_condition)

# After ToolNode runs, route back to "agent" so it can see the tool results
# and either call more tools or produce a final answer.
graph_builder.add_edge("tools", "agent")

graph = graph_builder.compile(checkpointer=MemorySaver())

def main() -> None:
    config = {"configurable": {"thread_id": "demo-tool-agent"}}

    print("--- Order status lookup ---")
    # Expected flow: agent calls get_order_status("A123") → ToolNode → agent answers.
    result = graph.invoke(
        {"messages": [("user", "What's the status of order A123?")], "tool_call_count": 0},
        config
    )
    print(result["messages"][-1].content)
    # Output:
    # Order A123 is currently **out for delivery**. It should arrive soon —
    # is there anything else I can help you with?

    print("\n--- Multi-step: lookup then refund ---")
    config2 = {"configurable": {"thread_id": "demo-tool-agent-2"}}
    # Expected flow: agent calls get_order_status("B456") → sees "delivered"
    # → calls issue_refund("B456", 19.50) → ToolNode → agent answers.
    result2 = graph.invoke(
        {
            "messages": [("user", "Check order B456's status, and if it's delivered, refund $19.50 for it.")],
            "tool_call_count": 0,
        },
        config2,
    )
    print(result2["messages"][-1].content)
    # Output:
    # Done! Here's a summary:
    # - **Order B456 status:** Delivered ✅
    # - **Refund issued:** $19.50
    # The refund has been processed successfully. Is there anything else
    # you'd like me to help with?

    print("\n--- Tool-level guardrail: refund amount exceeding order total ---")
    config3 = {"configurable": {"thread_id": "demo-tool-agent-3"}}
    # Expected flow: agent calls issue_refund("A123", 500) → ToolNode returns
    # soft-error string → agent relays the error to the user (no exception raised).
    result3 = graph.invoke(
        {"messages": [("user", "Refund $500 for order A123.")], "tool_call_count": 0},
        config3
    )
    print(result3["messages"][-1].content)
    # Output:
    # I wasn't able to process that refund. Here's what happened:
    # - **Order A123** exists and is currently **out for delivery**.
    # - However, the original order total was only **$49.99**, so a $500 refund
    #   isn't possible — refunds can't exceed the amount paid.
    # Would you like me to instead refund the full **$49.99** for order A123?

--- Order status lookup ---
Order A123 is currently **out for delivery**. It should arrive soon — is there
anything else I can help you with?

--- Multi-step: lookup then refund ---
Done! Here's a summary:

- **Order B456 status:** Delivered ✅
- **Refund issued:** $19.50

The refund has been processed successfully. Is there anything else you'd like me
to help with?

--- Tool-level guardrail: refund amount exceeding order total ---
I wasn't able to process that refund. Here's what happened:

- **Order A123** exists and is currently **out for delivery**.
- However, the original order total was only **$49.99**, so a $500 refund isn't
  possible — refunds can't exceed the amount paid.

Would you like me to instead refund the full **$49.99** for order A123? Or if you
had a different order in mind (perhaps a typo in the ID), let me know and I can
look it up.