import questionary
from langgraph.graph import StateGraph, START
from langgraph.types import interrupt, Command
from langgraph.checkpoint.memory import InMemorySaver
from typing import TypedDict, List

class GraphState(TypedDict):
    human_value: str | None
    foo: str | None

# Node that triggers interrupt

def interrupt_node(state: GraphState) -> GraphState:
    # Trigger interrupt with payload
    payload = {
        "type": "confirm",
        "question": "Уверены, что хотите продолжить?",
        "allow_responds": ["approve", "reject"],
    }
    # interrupt returns None; graph pauses until resumed
    interrupt(payload)
    # After resume, state will contain the answer in payload['answer']
    answer = state.get("answer")
    state["human_value"] = answer
    return state

builder = StateGraph(GraphState)
builder.add_node("interrupt", interrupt_node)
builder.set_entry_point("interrupt")
builder.add_edge(START, "interrupt")
builder.add_edge("interrupt", "interrupt")  # loop to finish
graph = builder.compile(checkpointer=InMemorySaver())

# Run graph with interrupt handling
config = {"configurable": {"thread_id": "demo"}}

# Initial state
state: GraphState = {"human_value": None, "foo": None}

# Stream execution
for chunk in graph.stream(state, config):
    if "__interrupt__" in chunk:
        # Extract payload
        payload = chunk["__interrupt__"][0].value
        # Show question
        answer = questionary.select(
            payload["question"],
            choices=payload["allow_responds"],
        ).ask()
        # Resume with answer
        payload["answer"] = answer
        resume = Command(resume=payload)
        for _ in graph.stream(resume, config):
            pass
    else:
        # Final state output
        print("Final state:", chunk)