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, Dict, Any

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

# Node that triggers interrupt
async def interrupt_node(state: GraphState) -> GraphState:
    # Trigger interrupt with structured payload
    payload = {
        "type": "confirm",
        "question": "Уверены, что хотите продолжить?",
        "allow_responds": ["approve", "reject"],
    }
    # interrupt returns None; graph pauses until resumed
    interrupt(payload)
    # After resume, state will contain the payload with added answer
    # Extract answer and store in state
    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.add_edge(START, "interrupt_node")
# No further nodes; graph ends after node
builder.set_finish_point("interrupt_node")

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

# Run graph with interrupt handling
config = {"configurable": {"thread_id": "thread-1"}}
stream = graph.stream({}, config)

for chunk in stream:
    if "__interrupt__" in chunk:
        # Get the interrupt payload
        interrupt_payload = chunk["__interrupt__"][0].value
        # Show question to user
        answer = questionary.select(
            interrupt_payload["question"],
            choices=interrupt_payload["allow_responds"],
        ).ask()
        # Add answer to payload and resume
        interrupt_payload["answer"] = answer
        stream = graph.stream(Command(resume=interrupt_payload), config)
        continue
    # Print final state when finished
    if "node" in chunk:
        print("Final state:", chunk["node"]["human_value"])
        break