"""
LangGraph Human‑in‑the‑Loop demo.

This script demonstrates how to use LangGraph’s interrupt / resume
mechanism to pause a graph, ask the user for a decision, and then
resume execution.

Requirements
------------
* langgraph
* questionary

Run the script with:

    pip install -r requirements.txt
    python main.py

The script will pause at the interrupt node, display a question in the
terminal, and wait for your response. After you answer, the graph will
continue and print the final state.
"""

from typing import TypedDict

from langgraph.graph import StateGraph, START
from langgraph.constants import interrupt
from langgraph.types import Command
from langgraph.checkpoint.memory import InMemorySaver
import questionary

# ---------------------------------------------------------------------------
# 1. State definition
# ---------------------------------------------------------------------------
class GraphState(TypedDict):
    """State of the graph.

    * ``human_value`` – value supplied by the user via interrupt.
    * ``foo`` – example of an initial value that could be used by the graph.
    """
    human_value: str
    foo: str

# ---------------------------------------------------------------------------
# 2. Node that triggers an interrupt
# ---------------------------------------------------------------------------
def interrupt_node(state: GraphState) -> GraphState:
    """Node that pauses the graph and asks the user for confirmation.

    The node calls :func:`langgraph.constants.interrupt` with a payload
    containing the question and allowed answers.  Execution will pause
    until the graph is resumed with a ``Command`` that contains the
    user’s response.
    """
    # Payload that will be sent to the interrupt handler.
    interrupt_payload = {
        "type": "confirm",
        "question": "Do you want to continue?",
        "allow_responds": ["approve", "reject"],
    }
    # Trigger the interrupt – execution stops here until resumed.
    interrupt(interrupt_payload)
    # After the graph is resumed, the same payload will be returned
    # to this node via the ``resume`` argument.  The node simply
    # returns the (possibly updated) state.
    return state

# ---------------------------------------------------------------------------
# 3. Build the graph
# ---------------------------------------------------------------------------
# Create a graph with an in‑memory checkpoint so we can resume after the
# interrupt.
graph = StateGraph(GraphState)
graph.add_node("interrupt_node", interrupt_node)
# The graph has only one node – it is both the entry and the finish point.
graph.set_entry_point("interrupt_node")
graph.set_finish_point("interrupt_node")
# Compile the graph with a checkpoint so the state is preserved across
# the pause/resume cycle.
graph.compile(checkpointer=InMemorySaver())

# ---------------------------------------------------------------------------
# 4. Run the graph with interrupt handling
# ---------------------------------------------------------------------------

def run_graph() -> None:
    """Execute the graph, handling interrupts in a loop.

    The function starts the graph, then iterates over the stream of
    chunks.  When an interrupt chunk is encountered, the question is
    displayed using ``questionary``.  The user’s answer is added to the
    payload and the graph is resumed with a ``Command``.
    """
    thread_id = "demo_thread"
    config = {"configurable": {"thread_id": thread_id}}

    # Start the stream.  ``{}`` is the initial state – the graph will
    # fill in the missing fields.
    stream = graph.stream({}, config)

    while True:
        try:
            chunk = next(stream)
        except StopIteration:
            # The stream has finished.
            break

        # -------------------------------------------------------------------
        # Handle interrupt chunks
        # -------------------------------------------------------------------
        if "__interrupt__" in chunk:
            # The interrupt payload is the first element of the list.
            interrupt_payload = chunk["__interrupt__"][0].value
            # Show the question and get the user’s answer.
            answer = questionary.select(
                interrupt_payload["question"],
                choices=interrupt_payload["allow_responds"],
            ).ask()
            # Attach the answer to the payload and resume.
            interrupt_payload["answer"] = answer
            # Resume the graph – this will produce a new stream.
            stream = graph.stream(Command(resume=interrupt_payload), config)
            continue

        # -------------------------------------------------------------------
        # Normal chunks – just print them.
        # -------------------------------------------------------------------
        print(chunk)

    # After the stream ends, fetch the final state from the checkpoint.
    final_state = graph.checkpointer.get_state(thread_id)
    print("\nFinal state:", final_state)

if __name__ == "__main__":
    run_graph()