# solution.py

"""
Мини‑игра «выбери свою историю» с прерыванием.
Использует:
- LangGraph для построения графа и управления прерываниями
- LangChain + langchain-openai для вызова LLM (ChatOpenAI)
- questionary для интерактивного выбора в консоли
"""

# ────────────────────── Импорты ──────────────────────

from typing import TypedDict, List, Dict, Any

import questionary  # pip install questionary

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

# LangChain + OpenAI LLM
from langchain_openai import ChatOpenAI

# ────────────────────── Состояние графа ──────────────────────

class State(TypedDict):
    """Состояние графа."""
    # Текст истории, генерируемый LLM
    story: str
    # Ответ пользователя после прерывания
    human_value: str | None


# ────────────────────── Узел с прерыванием ──────────────────────

def interrupt_node(state: State) -> State:
    """
    Генерирует часть истории и ставит граф на паузу,
    ожидая выбора игрока.
    """
    # Запрос к LLM для создания начала истории
    llm = ChatOpenAI(model="gpt-4o-mini")  # можно заменить на любой доступный LLM
    prompt = (
        "Создай короткую интригующую историю про героя, "
        "который стоит перед выбором. В конце предложи три варианта действий."
    )
    response = llm.invoke(prompt)
    state["story"] = response.content

    # Прерывание: пользователь выбирает вариант
    interrupt_payload = {
        "type": "alert",
        "question": "Выберите действие героя:",
        "allow_responds": ["Выйти", "Бороться", "Спрятаться"],
    }
    # Возврат прерывания; выполнение не продолжится до возобновления
    return interrupt(interrupt_payload)


# ────────────────────── Граф ──────────────────────

graph_builder = StateGraph(State)

# Добавляем узел и переход из START в него
graph_builder.add_node("interrupt", interrupt_node)
graph_builder.set_entry_point(START)
graph_builder.add_edge(START, "interrupt")
# После прерывания возвращаемся к тому же узлу для завершения
graph_builder.add_edge("interrupt", "interrupt")

# Чекпоинтер (InMemorySaver) сохраняет состояние между паузами
checkpoint = InMemorySaver()
graph = graph_builder.compile(checkpointer=checkpoint)


# ────────────────────── Запуск и обработка прерываний ──────────────────────

def main() -> None:
    """
    Основной цикл запуска графа с обработкой интерактивных прерываний.
    """
    # Уникальный идентификатор потока (можно использовать любой строковый ID)
    thread_id = "interactive_story"

    # Инициализируем состояние
    init_state: State = {"story": "", "human_value": None}

    # Запускаем поток генерации
    stream = graph.stream(
        Command(resume=init_state),
        configurable={"thread_id": thread_id},
    )

    for chunk in stream:
        # Если в чанке есть прерывание, обрабатываем его
        if "__interrupt__" in chunk:
            interrupt_info = chunk["__interrupt__"][0].value  # dict с payload
            rprint(f"\n=== Прервано: {interrupt_info['type']} ===")
            rprint(interrupt_info["question"])
            # Выводим варианты и получаем выбор пользователя
            answer = questionary.select(
                interrupt_info["question"],
                choices=interrupt_info["allow_responds"],
            ).ask()
            if answer is None:
                raise RuntimeError("Пользователь отменил ввод.")
            rprint(f"\nВы выбрали: {answer}")

            # Добавляем ответ в payload и возобновляем граф
            interrupt_info["human_value"] = answer
            stream = graph.stream(
                Command(resume=interrupt_info),
                configurable={"thread_id": thread_id},
            )
            continue

        # Выводим обычные сообщения от LLM (если есть)
        if "story" in chunk:
            rprint(f"\nИстория: {chunk['story']}")
        if "human_value" in chunk and chunk["human_value"]:
            rprint(f"\nОтвет игрока сохранён в состоянии: {chunk['human_value']}")

    # После завершения выводим итоговое состояние
    final_state = graph.get_state(configurable={"thread_id": thread_id})
    rprint("\n=== Итоговое состояние ===")
    rprint(final_state)


if __name__ == "__main__":
    main()