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

class GraphState(TypedDict):
    messages: Annotated[list, add_messages]
    human_value: str | None

async def interrupt_node(state: GraphState) -> GraphState:
    # Trigger interrupt asking user to confirm
    payload = {
        "type": "confirm",
        "question": "Do you want to continue?",
        "options": ["yes", "no"],
    }
    # interrupt returns None, graph pauses
    interrupt(payload)
    # After resume, payload will have 'answer'
    answer = state.get("human_value")
    return {"messages": state["messages"], "human_value": answer}

builder = StateGraph(GraphState)
builder.add_node("interrupt_node", interrupt_node)
builder.add_edge(START, "interrupt_node")
builder.add_edge("interrupt_node", START)
graph = builder.compile(checkpointer=InMemorySaver())

async def main():
    config = {"configurable": {"thread_id": "session-1"}}
    # Start stream
    async for chunk in graph.stream({"messages": [], "human_value": None}, config):
        if "__interrupt__" in chunk:
            payload = chunk["__interrupt__"][0].value
            answer = questionary.select(payload["question"], choices=payload["options"]).ask()
            # resume with answer
            payload["answer"] = answer
            # store answer in state
            chunk["human_value"] = answer
            await graph.stream(Command(resume=payload), config)
        else:
            print(chunk)

if __name__ == "__main__":
    import asyncio
    asyncio.run(main())