import sys
from typing import TypedDict

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


# ---------- 1. Состояние графа ----------
class GraphState(TypedDict):
    human_value: str | None
    foo: str | None


# ---------- 2. Узел с прерыванием ----------
def node_with_interrupt(state: GraphState) -> GraphState:
    # Создаём объект прерывания
    interrupt_payload = {
        "type": "confirm",
        "question": "Уверены, что хотите продолжить?",
        "allow_responds": ["approve", "reject"],
    }
    # Вызываем прерывание – выполнение остановится до возобновления.
    # После вызова graph.stream(Command(resume=answer), config)
    # interrupt() вернёт значение, переданное в resume=.
    answer = interrupt(interrupt_payload)

    # Сохраняем ответ пользователя в состоянии
    return {"human_value": answer, "foo": state.get("foo")}


# ---------- 3. Сборка графа ----------
builder = StateGraph(GraphState)
builder.add_node("interrupt_node", node_with_interrupt)
builder.add_edge(START, "interrupt_node")
graph = builder.compile(checkpointer=InMemorySaver())


# ---------- 4. Цикл запуска с обработкой прерывания ----------
def run_graph() -> None:
    init_state: GraphState = {"human_value": None, "foo": None}
    config = {"configurable": {"thread_id": "demo_thread"}}

    # Первый запуск – граф дойдёт до interrupt() и остановится
    print("[Graph] Starting new run...")
    for chunk in graph.stream(init_state, config):
        if "__interrupt__" in chunk:
            # Получаем объект прерывания
            interrupt_obj = chunk["__interrupt__"][0].value

            print("\n=== Появилось прерывание ===")
            print(f"Тип: {interrupt_obj['type']}")
            print(f"Вопрос: {interrupt_obj['question']}")
            print(f"Варианты: {', '.join(interrupt_obj['allow_responds'])}")

            # Запрашиваем ответ пользователя через questionary
            answer = questionary.select(
                interrupt_obj["question"],
                choices=interrupt_obj["allow_responds"],
            ).ask()

            if answer is None:
                sys.exit("Отмена пользователем")

            # Возобновляем граф, передавая ответ напрямую в Command(resume=)
            print("\n[Graph] Resuming after interrupt...")
            for resumed_chunk in graph.stream(Command(resume=answer), config):
                print(resumed_chunk)

    # После завершения выводим финальное состояние через get_state
    final_state = graph.get_state(config)
    print("\n=== Финальное состояние ===")
    print(final_state.values)


if __name__ == "__main__":
    run_graph()