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

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

# Node that triggers interrupt

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

# Build graph
builder = StateGraph(GraphState)
builder.add_node("interrupt_node", interrupt_node)
builder.set_entry_point("interrupt_node")
builder.set_finish_point("interrupt_node")
# Use InMemorySaver for checkpointing
checkpoint = InMemorySaver()
graph = builder.compile(checkpointer=checkpoint)

# Run graph with interrupt handling
config = {"configurable": {"thread_id": "thread-1"}}
# Initial state
state: GraphState = {"human_value": None, "foo": None}

# Start streaming
stream = graph.stream(state, config)
for chunk in stream:
    if "__interrupt__" in chunk:
        # Extract payload
        payload = chunk["__interrupt__"][0].value
        # Show question to user
        answer = questionary.select(
            payload["question"],
            choices=payload["allow_responds"],
        ).ask()
        # Add answer to payload and resume
        payload["answer"] = answer
        # Resume graph
        stream = graph.stream(Command(resume=payload), config)
        # Continue processing resumed stream
        for subchunk in stream:
            if "__interrupt__" in subchunk:
                # Should not happen in this simple example
                continue
            print(subchunk)
    else:
        print(chunk)

print("Final state:", state)