#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Human‑in‑the‑loop demo for LangGraph.
Run with:
    python interrupt_demo.py
"""

from __future__ import annotations

import uuid
from typing import Optional

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


# ------------------------------------------------------------------
# 1️⃣ Состояние графа
# ------------------------------------------------------------------
class State(TypedDict):
    """Структура состояния LangGraph."""

    foo: str                     # начальные данные (можно использовать как угодно)
    human_value: Optional[str]   # будет заполнено после пользовательского ответа


# ------------------------------------------------------------------
# 2️⃣ Узел с прерыванием
# ------------------------------------------------------------------
def node(state: State) -> dict:
    """
    Узел, который останавливает выполнение и запрашивает у пользователя подтверждение.
    После возобновления он сохраняет ответ в `human_value`.
    """
    # 1. Отправляем запрос на прерывание
    interrupt_payload = interrupt(
        {
            "type": "confirm",
            "question": "Уверены, что хотите продолжить?",
            "allow_responds": ["approve", "reject"],
        }
    )

    # 2. После возобновления `interrupt_payload` будет содержать поле `answer`
    answer = interrupt_payload["answer"]
    print(f"> Received an input from the interrupt: {answer}")

    # 3. Возвращаем обновлённое состояние
    return {"human_value": answer}


# ------------------------------------------------------------------
# 3️⃣ Сборка графа
# ------------------------------------------------------------------
builder = StateGraph(State)
builder.add_node("node", node)
builder.add_edge(START, "node")            # единственный узел

checkpointer = InMemorySaver()              # в памяти (для простоты)
graph = builder.compile(checkpointer=checkpointer)


# ------------------------------------------------------------------
# 4️⃣ Запуск и обработка прерываний
# ------------------------------------------------------------------
def main() -> None:
    config = {
        "configurable": {"thread_id": uuid.uuid4()},
    }

    # Инициализируем поток с начальным состоянием
    initial_state = {"foo": "some_initial_value"}

    for chunk in graph.stream(initial_state, config):
        # Если в чанке есть прерывание – обрабатываем его
        if "__interrupt__" in chunk:
            # `chunk["__interrupt__"]` – список объектов. Берём первый.
            interrupt_obj = chunk["__interrupt__"][0].value

            print("\n⚠️  Произошла остановка ⚠️")
            print(interrupt_obj)

            # Варианты ответа
            answer = questionary.select(
                interrupt_obj["question"],
                choices=interrupt_obj["allow_responds"],
            ).ask()

            # Добавляем ответ в объект прерывания и возобновляем граф
            interrupt_obj["answer"] = answer
            command = Command(resume=interrupt_obj)

            # Продолжаем поток после резюме
            for resumed_chunk in graph.stream(command, config):
                print(resumed_chunk)
        else:
            # Печатаем обычные чанки (состояния, сообщения и т.п.)
            print(chunk)


if __name__ == "__main__":
    main()