import questionary
from typing import TypedDict, Optional

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

# Define state schema
class State(TypedDict, total=False):
    human_value: Optional[str]
    foo: Optional[str]

# Node that triggers interrupt
async def interrupt_node(state: State) -> State:
    payload = {
        "type": "choice",
        "question": "Уверены, что хотите продолжить?",
        "options": ["approve", "reject"],
    }
    # Trigger interrupt
    return interrupt(payload)

# Build graph
graph_builder = StateGraph(State)
graph_builder.add_node("interrupt", interrupt_node)
graph_builder.set_entry_point("interrupt")
graph_builder.add_edge("interrupt", END)

# Compile graph with in-memory checkpointing
checkpoint = InMemorySaver()
graph = graph_builder.compile(checkpointer=checkpoint)

if __name__ == "__main__":
    initial_state: State = {"human_value": None, "foo": None}
    thread_id = "main-thread"

    # Run graph to detect interrupt
    for chunk in graph.stream(initial_state, thread_id=thread_id):
        if "__interrupt__" in chunk:
            payload = chunk["__interrupt__"]
            try:
                answer = questionary.select(
                    payload["question"], choices=payload["options"]
                ).ask()
            except Exception as e:
                print(f"Error during user input: {e}")
                answer = None
            # Store answer in state
            if answer is not None:
                initial_state["human_value"] = answer
            # Resume graph
            final_state = graph.invoke(Command(resume=payload), thread_id=thread_id)
            print("Final state:", final_state)
            break
    else:
        # No interrupt, graph finished normally
        print("Graph finished without interrupt.")