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

import uuid
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 State(TypedDict, total=False):
    foo: str
    human_value: str


def human_node(state: State) -> dict:
    payload = {
        "type": "alert",
        "question": "Уверены, что хотите продолжить?",
        "allow_responds": ["approve", "reject"],
    }
    resumed = interrupt(payload)

    if isinstance(resumed, dict):
        answer = resumed.get("answer", "")
    else:
        answer = str(resumed)

    print(f"> Received an input from the interrupt: {answer}")
    return {
        "foo": state.get("foo", ""),
        "human_value": answer,
    }


def build_graph():
    builder = StateGraph(State)
    builder.add_node("node", human_node)
    builder.add_edge(START, "node")
    return builder.compile(checkpointer=InMemorySaver())


def handle_interrupt(interrupts: tuple) -> dict:
    first = interrupts[0]
    payload = dict(first.value if hasattr(first, "value") else first)

    print("Произошла остановка")
    print(payload)
    print(f"!!! {payload.get('type', 'alert')} !!!")

    choice = questionary.select(
        payload["question"],
        choices=payload["allow_responds"],
    ).ask()

    payload["answer"] = choice or payload["allow_responds"][0]
    return payload


def run_hitl() -> None:
    graph = build_graph()
    config = {"configurable": {"thread_id": str(uuid.uuid4())}}
    stream_input: dict | Command = {"foo": "abc"}

    while True:
        interrupted = False
        for chunk in graph.stream(stream_input, config=config):
            if "__interrupt__" in chunk:
                payload = handle_interrupt(chunk["__interrupt__"])
                stream_input = Command(resume=payload)
                interrupted = True
                break
            print(chunk)

        if not interrupted:
            break

        for chunk in graph.stream(stream_input, config=config):
            print(chunk)


if __name__ == "__main__":
    run_hitl()