import questionary
from typing import TypedDict, List, Any, Optional
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.types import interrupt, Command

# 1. Описание состояния графа
class GraphState(TypedDict):
    foo: str
    human_value: Optional[str]
    status: str

# 2. Узел с прерыванием
def human_approval_node(state: GraphState) -> GraphState:
    """
    Узел, который ставит выполнение на паузу и ждет ответа пользователя.
    При вызове interrupt() граф останавливается. 
    Когда вызывается Command(resume=...), значение из resume попадает в результат interrupt().
    """
    # Формируем payload для прерывания
    interrupt_payload = {
        "type": "confirm",
        "question": "Уверены, что хотите продолжить?",
        "options": ["approve", "reject"]
    }
    
    # Вызываем interrupt. 
    # Если граф возобновлен через Command(resume=...), 
    # переменная 'user_response' получит значение из resume.
    user_response = interrupt(interrupt_payload)
    
    # Возвращаем обновленное состояние
    return {
        "human_value": user_response,
        "status": "completed"
    }

# 3. Сборка графа
def create_graph():
    workflow = StateGraph(GraphState)
    
    workflow.add_node("approval_node", human_approval_node)
    workflow.add_edge(START, "approval_node")
    workflow.add_edge("approval_node", END)
    
    # Используем InMemorySaver для поддержки чекпоинтов (необходим для прерываний)
    memory = InMemorySaver()
    return workflow.compile(checkpointer=memory)

def run_workflow():
    graph = create_graph()
    config = {"configurable": {"thread_id": "123"}}
    initial_state = {"foo": "initial_data", "human_value": None, "status": "running"}

    # Запуск графа
    print("--- Starting Graph Execution ---")
    
    # Первый запуск (до прерывания)
    # Используем stream_mode="updates" или "values". 
    # При interrupt в потоке появится специальный элемент.
    events = graph.stream(initial_command_or_state(initial_state), config, stream_mode="updates")
    
    final_state = None
    
    # Итерируемся по событиям
    for event in events:
        # Проверяем наличие прерывания в потоке через ключ __interrupt__
        if "__interrupt__" in event:
            # Извлекаем payload из прерывания
            # В LangGraph interrupt создает элемент в списке interrupts задачи
            interrupt_info = event["__interrupt__"][0].value
            
            print("\n--- Interrupt Detected ---")
            print(interrupt_info)
            print(f"!!! {interrupt_info['type'].upper()} !!!")
            
            # 4. Сбор ответа пользователя через questionary
            user_choice = questionary.select(
                interrupt_info["question"],
                choices=interrupt_info["options"]
            ).ask()
            
            print(f"> Received an input from the interrupt: {user_choice}")
            
            # 5. Возобновление графа через Command(resume=...)
            # Передаем выбранный ответ обратно в граф
            resume_events = graph.stream(
                Command(resume=user_choice),
                config,
                stream_mode="updates"
            )
            
            for resume_event in resume_events:
                final_state = resume_event
        else:
            # Обычное обновление узлов
            final_state = event

    # Вывод итогового состояния (из последнего события или snapshot)
    print("\n--- Final State ---")
    snapshot = graph.get_state(config)
    print(snapshot.values)

def initial_command_or_state(initial_state):
    """Вспомогательная функция для корректного старта."""
    return initial_state

if __name__ == "__main__":
    run_workflow()