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

# State definition
class GraphState(TypedDict):
    human_value: Annotated[str | None, "value chosen by user"]
    foo: Annotated[str, "initial data placeholder"]

# Node that triggers interrupt

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

# Build graph
builder = StateGraph(GraphState)
builder.add_node("interrupt_node", interrupt_node)
builder.add_edge(START, "interrupt_node")
builder.add_edge("interrupt_node", END)

# Compile with checkpoint
graph = builder.compile(checkpointer=InMemorySaver())

# Run graph with stream and handle interrupt
config = {"configurable": {"thread_id": "demo"}}

# First run: will pause at interrupt
stream = graph.stream({}, config)
for chunk in stream:
    if "__interrupt__" in chunk:
        # Extract payload
        payload = chunk["__interrupt__"][0].value
        # Show question and get answer
        answer = questionary.select(
            payload["question"],
            choices=payload["allow_responds"],
        ).ask()
        # Add answer to payload
        payload["answer"] = answer
        # Resume graph
        stream = graph.stream(Command(resume=payload), config)
        continue
    # Print normal output
    print(chunk)

# After stream ends, print final state
print("Final state:", stream.final_state)