from typing import TypedDict, Dict, Any

from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import InMemorySaver


class GraphState(TypedDict):
    human_value: str | None
    __interrupt__: dict[str, Any] | None  # for interrupt payload
    __resume__: str | None  # for resume payload


def interrupt_node(state: GraphState) -> Dict[str, Any]:
    """Trigger an interrupt asking the user to choose a value."""
    return {
        "__interrupt__": {
            "type": "question",
            "question": "Choose a value:",
            "options": ["Option A", "Option B", "Option C"],
        }
    }


def resume_node(state: GraphState) -> Dict[str, Any]:
    """Store the user's choice and finish."""
    # The chosen option is passed via the __resume__ key
    if state.get("__resume__") is not None:
        state["human_value"] = state["__resume__"]
    return {"__end__": True}


builder = StateGraph(GraphState)
builder.add_node("interrupt", interrupt_node)
builder.add_node("resume", resume_node)

# The graph starts with the interrupt node
builder.set_entry_point("interrupt")

# After a successful resume, we end the graph
builder.add_edge(START, "interrupt")
builder.add_conditional_edges(
    "interrupt",
    lambda x: "__interrupt__" in x,
    {"__interrupt__": "resume"},
)
builder.add_edge("resume", END)

graph = builder.compile(checkpointer=InMemorySaver())

# Run the graph and handle interrupts
state: GraphState = {"human_value": None, "__interrupt__": None, "__resume__": None}
while True:
    result = graph.invoke(state)
    if "__interrupt__" in result and result["__interrupt__"] is not None:
        interrupt_info = result["__interrupt__"]
        print(interrupt_info["question"])
        for idx, opt in enumerate(interrupt_info["options"], 1):
            print(f"{idx}. {opt}")
        choice_idx = int(input("Enter number: ")) - 1
        chosen = interrupt_info["options"][choice_idx]
        # Resume the graph with the chosen value
        state = graph.invoke(state, {"__resume__": chosen})
    else:
        # Graph finished
        break

print("Final state:", state)