import os
from typing import TypedDict, List, Optional

import questionary
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver

# ------------------------------------------------------------------
# 1. Определяем состояние графа
# ------------------------------------------------------------------
class StoryState(TypedDict):
    # Тема/завязка истории
    theme: Optional[str]
    # Текст сцены, сгенерированный LLM
    scene: Optional[str]
    # Список вариантов выбора
    options: List[str]
    # Выбранный пользователем вариант
    choice: Optional[str]
    # Концовка, сгенерированная LLM
    ending: Optional[str]


# ------------------------------------------------------------------
# 2. Инициализируем LLM
# ------------------------------------------------------------------
# Убедитесь, что переменная окружения OPENAI_API_KEY установлена
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7)


# ------------------------------------------------------------------
# 3. Узлы графа
# ------------------------------------------------------------------
def generate_scene(state: StoryState) -> StoryState:
    """
    Генерируем начальную сцену и варианты выбора.
    """
    theme = state.get("theme") or "Вы просыпаетесь в загадочном лесу."
    prompt = f"""
    Напиши короткую сцену (не более 150 слов) для истории, где главный герой {theme}
    и предложи три варианта действий, которые может предпринять герой.
    Выведи варианты в виде списка, каждый с номером 1-3.
    """
    response = llm.invoke(prompt)
    # Разделяем на сцену и варианты
    lines = response.content.strip().split("\n")
    scene_lines = []
    options = []
    for line in lines:
        if line.strip().startswith(("1.", "2.", "3.")):
            options.append(line.strip())
        else:
            scene_lines.append(line.strip())
    state["scene"] = " ".join(scene_lines)
    state["options"] = options
    return state


def interrupt_choice(state: StoryState) -> StoryState:
    """
    Пауза и запрос выбора у пользователя.
    """
    print("\n--- Сцена ---")
    print(state["scene"])
    print("\n--- Выберите действие ---")
    for opt in state["options"]:
        print(opt)
    choice = questionary.select(
        "Выберите вариант:",
        choices=[opt.split(".", 1)[1].strip() for opt in state["options"]],
    ).ask()
    state["choice"] = choice
    return state


def generate_ending(state: StoryState) -> StoryState:
    """
    Генерируем короткую концовку на основе выбранного варианта.
    """
    prompt = f"""
    Герой выбрал: "{state['choice']}".
    Напиши короткую концовку (не более 100 слов), которая логично следует из выбранного действия.
    """
    response = llm.invoke(prompt)
    state["ending"] = response.content.strip()
    return state


# ------------------------------------------------------------------
# 4. Создаём граф
# ------------------------------------------------------------------
builder = StateGraph(StoryState)

# Добавляем узлы
builder.add_node("scene", generate_scene)
builder.add_node("choice", interrupt_choice)
builder.add_node("ending", generate_ending)

# Определяем переходы
builder.set_entry_point("scene")
builder.add_edge("scene", "choice")
builder.add_edge("choice", "ending")
builder.add_edge("ending", END)

# Сохраняем состояние в памяти (необязательно, но удобно)
checkpoint = MemorySaver()
graph = builder.compile(checkpointer=checkpoint)


# ------------------------------------------------------------------
# 5. Запуск
# ------------------------------------------------------------------
def main():
    # Инициализируем состояние
    init_state: StoryState = {
        "theme": None,
        "scene": None,
        "options": [],
        "choice": None,
        "ending": None,
    }

    # Запускаем граф
    result = graph.invoke(init_state)

    # Выводим финальный результат
    print("\n=== Концовка ===")
    print(result["ending"])


if __name__ == "__main__":
    main()