import sys
from typing import TypedDict, List, Dict, Any

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

# Define state
class State(TypedDict):
    human_value: str | None
    foo: str | None

# Node that triggers interrupt
def interrupt_node(state: State) -> State:
    # Trigger interrupt with question and options
    interrupt_payload = {
        "type": "confirm",
        "question": "Уверены, что хотите продолжить?",
        "allow_responds": ["approve", "reject"],
    }
    # Raise interrupt; graph will pause until resumed
    interrupt(interrupt_payload)
    # After resume, the payload will be merged into state via resume
    # We expect state to contain 'human_value' set by resume
    return state

# Build graph
graph = StateGraph(State)
graph.add_node("interrupt", interrupt_node)
graph.set_entry_point(START)
graph.add_edge(START, "interrupt")
# No further nodes; graph ends after interrupt node
graph.set_finish_point("interrupt")

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

# Run graph with interrupt handling
if __name__ == "__main__":
    thread_id = "demo_thread"
    config = {"configurable": {"thread_id": thread_id}}
    # Start stream
    stream = compiled.stream({}, config)
    for chunk in stream:
        if "__interrupt__" in chunk:
            # Extract interrupt payload
            interrupt_obj = chunk["__interrupt__"][0]
            payload = interrupt_obj.value
            # Show question to user
            answer = questionary.select(
                payload["question"],
                choices=payload["allow_responds"],
            ).ask()
            # Resume with answer
            resume_payload = {"human_value": answer}
            stream = compiled.stream(Command(resume=resume_payload), config)
            for subchunk in stream:
                if "human_value" in subchunk:
                    print("\nFinal state:", subchunk)
                    sys.exit(0)
        else:
            # Normal output (none expected here)
            pass