"""Human-in-the-loop: кастомное прерывание (interrupt / resume) в LangGraph."""
from __future__ import annotations

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 Command, interrupt


class GraphState(TypedDict):
    """Состояние графа: начальные данные и ответ пользователя."""

    foo: str
    human_value: str


def human_node(state: GraphState) -> GraphState:
    """Узел с кастомным прерыванием — ждёт ответ пользователя."""
    payload = interrupt(
        {
            "type": "alert",
            "question": "Уверены что хотите продолжить?",
            "allow_responds": ["approve", "reject"],
        }
    )

    answer = payload.get("answer", "")
    print(f"!!! {payload.get('type', 'alert')} !!!")
    print(f"> Received an input from the interrupt: {answer}")

    return {
        "foo": state.get("foo", ""),
        "human_value": answer,
    }


def _ask_user(payload: dict) -> dict:
    """Показать вопрос и записать ответ в payload."""
    print(payload)
    options = payload.get("allow_responds") or ["approve", "reject"]
    question = payload.get("question", "Выберите вариант:")
    choice = questionary.select(question, choices=options).ask()
    if choice is None:
        choice = options[0]
    updated = dict(payload)
    updated["answer"] = choice
    return updated


def run_graph() -> GraphState:
    """Запуск графа с обработкой прерывания и возобновлением."""
    builder = StateGraph(GraphState)
    builder.add_node("human_node", human_node)
    builder.add_edge(START, "human_node")

    graph = builder.compile(checkpointer=InMemorySaver())
    config = {"configurable": {"thread_id": "hitl-interrupt-demo"}}
    initial: GraphState = {"foo": "initialized", "human_value": ""}

    final_state: GraphState = initial

    stream = graph.stream(initial, config)
    for chunk in stream:
        if "__interrupt__" in chunk:
            print("Произошла остановка")
            interrupt_tuple = chunk["__interrupt__"]
            payload = dict(interrupt_tuple[0].value)
            resumed_payload = _ask_user(payload)

            resume_stream = graph.stream(Command(resume=resumed_payload), config)
            for resume_chunk in resume_stream:
                if "human_node" in resume_chunk:
                    final_state = resume_chunk["human_node"]
                    print({"node": resume_chunk})
        elif "human_node" in chunk:
            final_state = chunk["human_node"]
            print(chunk)

    return final_state


def main() -> None:
    result = run_graph()
    print("\n=== Итоговое состояние ===")
    print(result)


if __name__ == "__main__":
    main()