from typing import TypedDict, List

import questionary
from langchain_openai import ChatOpenAI
from langchain.prompts import ChatPromptTemplate
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.constants import START
from langgraph.graph import StateGraph
from langgraph.types import interrupt, Command


# ---------- 1. Состояние графа ----------

class GraphState(TypedDict):
    """Состояние графа с отдельными полями для каждого этапа."""
    topic: str          # Тема (входной параметр)
    setup: str          # Сгенерированная завязка от LLM
    choices: List[str]  # Список вариантов выбора (3 штуки)
    human_choice: str   # Выбранный пользователем вариант
    ending: str         # Концовка от LLM


# ---------- LangChain: промпты и цепочки ----------

llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0.7)

# Цепочка для генерации завязки и вариантов
scene_prompt = ChatPromptTemplate.from_template(
    "Тема: {topic}.\n"
    "Придумай короткую завязку (2–3 предложения) и ровно 3 варианта поступка героя.\n"
    "Ответь строго в формате:\n"
    "ЗАВЯЗКА: <текст завязки>\n"
    "ВАРИАНТЫ:\n"
    "1) <первый вариант>\n"
    "2) <второй вариант>\n"
    "3) <третий вариант>"
)
scene_chain = scene_prompt | llm

# Цепочка для генерации концовки
ending_prompt = ChatPromptTemplate.from_template(
    "Завязка: {setup}\n"
    "Выбор пользователя: {human_choice}\n\n"
    "Допиши короткую концовку истории (2–3 предложения)."
)
ending_chain = ending_prompt | llm


# ---------- 2. Узел: генерация сцены + прерывание + концовка ----------

def story_node(state: GraphState) -> GraphState:
    """
    Шаг 1: генерирует завязку и варианты через LangChain-цепочку.
    Шаг 2: вызывает interrupt() — граф ставится на паузу.
    Шаг 3: после resume дописывает концовку через LangChain-цепочку.
    """

    # --- Шаг 1: генерируем завязку и варианты через LangChain ---
    scene_response = scene_chain.invoke({"topic": state["topic"]})
    raw = scene_response.content.strip()

    # Парсим ответ: разделяем завязку и варианты
    setup = ""
    choices = []

    lines = raw.splitlines()
    in_choices = False
    for line in lines:
        line = line.strip()
        if line.upper().startswith("ЗАВЯЗКА:"):
            setup = line[len("ЗАВЯЗКА:"):].strip()
        elif line.upper().startswith("ВАРИАНТЫ:"):
            in_choices = True
        elif in_choices and line:
            # Убираем нумерацию вида "1)", "2)", "3)" или "1.", "2.", "3."
            for prefix in ("1)", "2)", "3)", "1.", "2.", "3."):
                if line.startswith(prefix):
                    line = line[len(prefix):].strip()
                    break
            if line:
                choices.append(line)

    # Защита: если парсинг не дал нужных данных
    if not setup:
        setup = raw
    if len(choices) < 3:
        choices = choices + [f"Вариант {i+1}" for i in range(len(choices), 3)]

    state["setup"] = setup
    state["choices"] = choices

    # --- Шаг 2: прерывание ---
    interrupt_payload = {
        "type": "choice",
        "question": setup + "\n\nЧто делаем?",
        "allow_responds": choices,
    }

    # Выполнение останавливается здесь до Command(resume=...)
    resume_value = interrupt(interrupt_payload)

    # После resume в resume_value — словарь с полем "answer"
    human_choice = resume_value.get("answer", choices[0])
    state["human_choice"] = human_choice

    # --- Шаг 3: генерируем концовку через LangChain ---
    ending_response = ending_chain.invoke({
        "setup": setup,
        "human_choice": human_choice,
    })
    state["ending"] = ending_response.content.strip()

    return state


# ---------- 3. Сборка графа ----------

builder = StateGraph(GraphState)
builder.add_node("story", story_node)
builder.add_edge(START, "story")

graph = builder.compile(checkpointer=InMemorySaver())


# ---------- 4. Цикл запуска с обработкой прерывания ----------

def main() -> None:
    topic = input("Введите тему истории (например, «космический кот»): ").strip()
    if not topic:
        topic = "космический кот"

    print(f"\nТема: {topic}\n")

    thread_id = "story_thread_1"
    config = {"configurable": {"thread_id": thread_id}}

    initial_state: GraphState = {
        "topic": topic,
        "setup": "",
        "choices": [],
        "human_choice": "",
        "ending": "",
    }

    # --- Первый запуск: до прерывания ---
    for chunk in graph.stream(initial_state, config=config):
        if "__interrupt__" in chunk:
            interrupt_obj = chunk["__interrupt__"][0].value  # наш payload

            print(f"\n[LLM] {interrupt_obj['question']}\n")

            # Показываем варианты через questionary
            answer = questionary.select(
                "Выберите действие:",
                choices=interrupt_obj["allow_responds"],
            ).ask()

            if answer is None:
                raise SystemExit("Отмена пользователем.")

            # Добавляем ответ в payload и возобновляем граф
            interrupt_obj["answer"] = answer

            # --- Возобновление через Command(resume=...) ---
            for resumed_chunk in graph.stream(
                Command(resume=interrupt_obj), config=config
            ):
                # После resume граф дописывает концовку — просто ждём завершения
                pass

    # Получаем финальное состояние
    final_state = graph.get_state(config).values

    print("\n" + "=" * 50)
    print("=== Итоговая история ===")
    print("=" * 50)
    print(f"\nЗавязка:\n{final_state['setup']}")
    print(f"\nВыбор: {final_state['human_choice']}")
    print(f"\n[LLM] {final_state['ending']}")
    print("\n" + "=" * 50)


if __name__ == "__main__":
    main()