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'
    # We expect the resumed payload to be the same dict with added 'answer'
    # The node receives the resumed payload as its return value
    # So we just return the state with human_value set
    return {"human_value": payload.get("answer", None), "foo": state.get("foo", "")}

# 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 interrupt node
graph = builder.compile(checkpointer=InMemorySaver())

# Run graph with interrupt handling
if __name__ == "__main__":
    config = {"configurable": {"thread_id": "hitl-demo"}}
    # Initial state
    state = {"human_value": None, "foo": "initial"}
    # Stream execution
    for chunk in graph.stream(state, config):
        if "__interrupt__" in chunk:
            # Extract 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
            # Resume graph
            for resume_chunk in graph.stream(Command(resume=interrupt_payload), config):
                # Print final state when graph finishes
                if "node" in resume_chunk:
                    print("Final state:", resume_chunk["node"])  
        else:
            # Normal output (none expected here)
            pass