START → classify_intent → draft_response → [conditional] → human_review → END
                                                  ↓ (no approval needed)
                                           finalize_response → END

# `llm` is the main reasoning model — used for open-ended drafting where
# response quality matters (order_status answers, general questions).
llm = get_llm("OPENROUTER_MODEL", "z-ai/glm-5.3-flash")

# `fast_llm` is used only for classification — a bounded, structured task
# that doesn't need a powerful model. In production, point
# OPENROUTER_FAST_MODEL at a smaller / cheaper slug (e.g. google/gemini-flash)
# to reduce cost without sacrificing answer quality.
# Defaults to the same model as `llm` when OPENROUTER_FAST_MODEL is unset.
fast_llm = get_llm("OPENROUTER_FAST_MODEL", "z-ai/glm-5.3-flash")

from typing import Annotated, Literal
from typing_extensions import TypedDict
from langgraph.graph.message import add_messages

class SupportState(TypedDict):
    messages: Annotated[list, add_messages]   # append-only conversation history
    intent: str                                # "refund" | "order_status" | "general"
    requires_approval: bool                    # True only for "refund"; gates human_review
    draft_response: str                        # LLM's draft, stored before being promoted to messages
    classification_reasoning: str             # LLM's explanation of its classification

import os
import sqlite3
from dotenv import load_dotenv

load_dotenv()

from typing import Annotated, Literal
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langgraph.checkpoint.sqlite import SqliteSaver  # file-backed checkpointer
from langgraph.types import interrupt, Command       # interrupt() pauses; Command resumes
from langchain_openai import ChatOpenAI


def get_llm(model_env_var="OPENROUTER_MODEL", default_model="z-ai/glm-5.3-flash", **kwargs):
    return ChatOpenAI(
        model=os.environ.get(model_env_var, default_model),
        api_key=os.environ["OPENROUTER_API_KEY"],
        base_url="https://openrouter.ai/api/v1",
        default_headers={"HTTP-Referer": "https://example.com", "X-Title": "LangGraph Series Examples"},
        **kwargs,
    )


llm = get_llm("OPENROUTER_MODEL", "z-ai/glm-5.3-flash")
fast_llm = get_llm("OPENROUTER_FAST_MODEL", "z-ai/glm-5.3-flash")


class SupportState(TypedDict):
    messages: Annotated[list, add_messages]
    intent: str
    requires_approval: bool
    draft_response: str
    classification_reasoning: str


# --------------- Node 1: classify_intent ---------------

def classify_intent(state: SupportState) -> dict:
    """
    Classify the customer's message into one of three intent buckets.
    Uses structured output (JSON schema) so the model MUST return a valid
    classification — no regex or free-text parsing needed.

    Returns:
        {
            "intent": "refund" | "order_status" | "general",
            "requires_approval": bool,       # True only for "refund"
            "classification_reasoning": str  # LLM's explanation of its choice
        }
    """
    user_msg = state["messages"][-1].content

    intent_llm = fast_llm.with_structured_output(
        {
            "title": "IntentClassification",
            "type": "object",
            "properties": {
                "intent": {"enum": ["refund", "order_status", "general"]},
                "reasoning": {"type": "string"},
            },
            "required": ["intent", "reasoning"],
        }
    )
    result = intent_llm.invoke(f"Classify this support request and briefly explain why: {user_msg}")
    return {
        "intent": result["intent"],
        # Refund requests always require human approval before a response is sent.
        "requires_approval": result["intent"] == "refund",
        "classification_reasoning": result["reasoning"],
    }


# --------------- Node 2: draft_response ---------------

def draft_response(state: SupportState) -> dict:
    """
    Draft a plain-text reply using the main (stronger) LLM.

    The draft is stored in state["draft_response"] but NOT yet added to
    state["messages"] — it must first pass through human_review (for refunds)
    or finalize_response (for everything else).
    """
    response = llm.invoke(state["messages"])
    return {"draft_response": response.content}


# --------------- Node 3a: human_review (refund path) ---------------

def human_review(state: SupportState) -> dict:
    """
    Human-in-the-loop gate: PAUSES the graph and waits for a human decision.

    interrupt() serialises the graph checkpoint to SqliteSaver and suspends
    execution. graph.invoke() returns early (without a final assistant message).
    When a human later calls graph.invoke(Command(resume={...}), config) the
    graph resumes from this exact point.

    Interrupt payload sent to the caller:
        {"intent": str, "reasoning": str, "draft": str}

    Resume payload expected via Command(resume={...}):
        {"approved": bool}
    """
    decision = interrupt(
        {
            "intent": state["intent"],
            "reasoning": state["classification_reasoning"],
            "draft": state["draft_response"],
        }
    )
    if decision["approved"]:
        return {"messages": [("assistant", state["draft_response"])]}
    # Human rejected or edited — use their override if provided
    return {"messages": [("assistant", decision.get("edited_response", "A team member will follow up shortly."))]}


# --------------- Node 3b: finalize_response (non-refund path) ---------------

def finalize_response(state: SupportState) -> dict:
    """
    Promote the draft into conversation history for non-refund requests.

    Without this node, draft_response would write to state["draft_response"]
    but nothing would add it to state["messages"]. The caller would receive
    the user's own question as result["messages"][-1], not the LLM's answer.
    This is a correctness requirement, not an optional node.
    """
    return {"messages": [("assistant", state["draft_response"])]}


# --------------- Router ---------------

def route_after_draft(state: SupportState) -> Literal["human_review", "finalize_response"]:
    """
    Route refunds to human_review; everything else to finalize_response.
    """
    return "human_review" if state["requires_approval"] else "finalize_response"


# --------------- Graph assembly ---------------

builder = StateGraph(SupportState)
builder.add_node("classify_intent", classify_intent)
builder.add_node("draft_response", draft_response)
builder.add_node("human_review", human_review)
builder.add_node("finalize_response", finalize_response)

# Every run: classify first, then draft unconditionally.
builder.add_edge(START, "classify_intent")
builder.add_edge("classify_intent", "draft_response")

# After drafting: refunds → human review, everything else → finalize.
builder.add_conditional_edges("draft_response", route_after_draft)

builder.add_edge("finalize_response", END)
builder.add_edge("human_review", END)

# SqliteSaver: checkpoint survives process restarts.
# Swap for PostgresSaver.from_conn_string(...) in production.
_conn = sqlite3.connect("langgraph_part5_demo.db", check_same_thread=False)
checkpointer = SqliteSaver(_conn)
graph = builder.compile(checkpointer=checkpointer)

def run_case(thread_id: str, user_message: str) -> None:
    config = {"configurable": {"thread_id": thread_id}}
    result = graph.invoke({"messages": [("user", user_message)]}, config)

    state = graph.get_state(config)
    if state.next:  # graph is paused at human_review
        pending = state.tasks[0].interrupts[0].value
        print(f"  [PAUSED] intent={pending['intent']!r} reasoning={pending['reasoning']!r}")
        print(f"  [PAUSED] draft awaiting approval: {pending['draft']}")
        # Simulate supervisor approving
        result = graph.invoke(Command(resume={"approved": True}), config)
        print(f"  [RESUMED] final: {result['messages'][-1].content}")
    else:
        print(f"  final: {result['messages'][-1].content}")


def main() -> None:
    print("--- Case 1: order status (no approval needed) ---")
    run_case("case-1", "Where is my order A123?")
    # Output:
    # final: I don't have access to order systems or databases, so I'm unable
    # to look up the status of order A123 or any other orders.
    # To track your order, you could try:
    # - Checking your email for a confirmation message with a tracking link
    # - Logging into your account on the company's website
    # - Contacting the company's customer service directly

    print("\n--- Case 2: refund request (requires human approval) ---")
    run_case("case-2", "I'd like a refund for order A123, it arrived damaged.")
    # Output:
    # [PAUSED] intent='refund'
    # [PAUSED] draft awaiting approval:
    #   I'm an AI assistant — I don't have access to any store's order systems,
    #   so I'm not able to look up order A123 or process a refund directly.
    #   But I can point you in the right direction: [... steps to contact retailer,
    #   document damage, check return window ...]
    # [RESUMED] final: [same approved draft sent to customer]

    print("\n--- Case 3: general question (no approval needed) ---")
    run_case("case-3", "What's your return policy?")
    # Output:
    # final: I'm an AI assistant, so I don't have a return policy since I'm not
    # a physical product. Is there something specific I can help you with today?

--- Case 1: order status (no approval needed) ---
I don't have access to order systems or databases, so I'm unable to look
up the status of order A123 or any other orders.

To track your order, you could try:
- Checking your email for a confirmation message with a tracking link
- Logging into your account on the company's website
- Contacting the company's customer service directly

--- Case 2: refund request (requires human approval) ---
[PAUSED] intent='refund'
[PAUSED] draft awaiting approval:
  I'm an AI assistant — I don't have access to any store's order systems,
  so I'm not able to look up order A123 or process a refund directly.
  But I can point you in the right direction: [... steps to contact retailer,
  document damage, check return window — truncated for brevity]

[RESUMED] final: [same approved draft sent to customer]

--- Case 3: general question (no approval needed) ---
I'm an AI assistant, so I don't have a return policy since I'm not
a physical product. Is there something specific I can help you with today?