"""
Задание 4: Текстовая игра на основе LLM + interrupt

Интерактивная история «выбери свою историю»:
LLM придумывает завязку и варианты, граф ставит на паузу (interrupt),
пользователь выбирает вариант, LLM дописывает концовку.
"""

import os
import re
from typing import TypedDict

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import StateGraph, START, END
from langgraph.types import interrupt, Command

try:
    import questionary
    HAS_QUESTIONARY = True
except ImportError:
    HAS_QUESTIONARY = False


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

class GameState(TypedDict):
    topic: str
    setup: str
    choices: list[str]
    selected_choice: str
    ending: str


# ─── Инициализация LLM ───────────────────────────────────────────────────────

api_base = os.environ.get("OPENAI_API_BASE", "https://api.openai.com/v1")
api_key = os.environ.get("OPENAI_API_KEY", "")
model_name = os.environ.get("OPENAI_MODEL", "gpt-4o-mini")

llm = ChatOpenAI(
    model=model_name,
    openai_api_base=api_base,
    openai_api_key=api_key,
    temperature=0.7,
)


# ─── Узел: генерация завязки + прерывание + концовка ─────────────────────────

def story_node(state: GameState) -> dict:
    """Один узел: генерирует завязку, прерывается на выбор, генерирует концовку."""

    if state.get("ending"):
        return {}

    if state.get("selected_choice") and state.get("setup"):
        response = llm.invoke([
            HumanMessage(content=(
                f"Завязка истории:\n{state['setup']}\n\n"
                f"Выбор пользователя: {state['selected_choice']}\n\n"
                f"Допиши короткую концовку (2–3 предложения) на русском языке."
            ))
        ])
        return {"ending": response.content}

    prompt = (
        f"Тема: {state['topic']}.\n\n"
        "Придумай короткую завязку для интерактивной истории (2–3 предложения) "
        "и ровно 3 варианта поступка героя.\n\n"
        "Ответь СТРОГО в формате:\n"
        "ЗАВЯЗКА: <текст завязки>\n"
        "ВАРИАНТ 1: <текст первого варианта>\n"
        "ВАРИАНТ 2: <текст второго варианта>\n"
        "ВАРИАНТ 3: <текст третьего варианта>\n"
    )

    response = llm.invoke([HumanMessage(content=prompt)])
    text = response.content

    setup_match = re.search(r"ЗАВЯЗКА:\s*(.+?)(?=ВАРИАНТ 1:)", text, re.DOTALL)
    setup = setup_match.group(1).strip() if setup_match else text

    choices = []
    for i in range(1, 4):
        m = re.search(rf"ВАРИАНТ {i}:\s*(.+?)(?=ВАРИАНТ {i+1}:|$)", text, re.DOTALL)
        if m:
            choices.append(m.group(1).strip())

    if not choices:
        choices = ["Пойти налево", "Пойти направо", "Остаться на месте"]

    user_choice = interrupt({
        "type": "choice",
        "question": f"\n{setup}\n\nЧто делаем?",
        "choices": choices,
    })

    return {
        "setup": setup,
        "choices": choices,
        "selected_choice": user_choice,
    }


# ─── Сборка графа ────────────────────────────────────────────────────────────

def build_graph():
    graph = StateGraph(GameState)
    graph.add_node("story", story_node)
    graph.add_edge(START, "story")
    graph.add_edge("story", END)

    memory = MemorySaver()
    return graph.compile(checkpointer=memory)


# ─── Интерактивный запуск ────────────────────────────────────────────────────

def run_game(topic: str):
    """Запустить интерактивную игру с заданной темой."""
    graph = build_graph()
    config = {"configurable": {"thread_id": "game-1"}}

    print("=" * 60)
    print("  🎮 ИНТЕРАКТИВНАЯ ИСТОРИЯ")
    print("=" * 60)
    print(f"  Тема: {topic}")
    print("=" * 60)

    state = {"topic": topic, "setup": "", "choices": [], "selected_choice": "", "ending": ""}
    stream = graph.stream(state, config)

    while True:
        result = None
        for chunk in stream:
            result = chunk

            if "__interrupt__" in chunk:
                intr = chunk["__interrupt__"][0]
                payload = intr.value

                print(f"\n{payload['question']}\n")
                for i, choice in enumerate(payload["choices"], 1):
                    print(f"  {i}. {choice}")

                if HAS_QUESTIONARY:
                    selected = questionary.select(
                        "Выберите вариант:",
                        choices=payload["choices"]
                    ).ask()
                else:
                    while True:
                        try:
                            idx = int(input("\nВведите номер варианта: ")) - 1
                            if 0 <= idx < len(payload["choices"]):
                                selected = payload["choices"][idx]
                                break
                        except (ValueError, IndexError):
                            pass
                        print("Неверный ввод, попробуйте снова.")

                stream = graph.stream(Command(resume=selected), config)
                break
        else:
            if result and "story" in result:
                story_state = result["story"]
                if story_state.get("ending"):
                    print(f"\n{'=' * 60}")
                    print(f"  📖 КОНЦОВКА")
                    print(f"{'=' * 60}")
                    print(f"\n{story_state['ending']}\n")
                    print(f"{'=' * 60}")
                    print("  🏁 КОНЕЦ ИСТОРИИ")
                    print(f"{'=' * 60}")
                    break
            elif result is None:
                break


if __name__ == "__main__":
    import sys
    topic = sys.argv[1] if len(sys.argv) > 1 else "космический кот"
    run_game(topic)