from typing import TypedDict
import json

import questionary
from langgraph.graph import StateGraph
from langgraph.constants import START
from langgraph.types import interrupt, Command
from langgraph.checkpoint.memory import InMemorySaver
from langchain_openai import ChatOpenAI
from rich.console import Console


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

class State(TypedDict):
    """Состояние графа."""
    story: str                # Текст истории, генерируемый LLM
    options: list[str]        # Варианты действий
    choice: str | None        # Выбор пользователя
    ending: str | None        # Концовка


# ────────────────────── Инструмент LLM ──────────────────────

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7)
console = Console()


def generate_story(state: State) -> State:
    """Генерирует начало истории и варианты действий."""
    prompt = (
        "Generate a short adventure intro and three choices in JSON format.\n"
        "The JSON should have keys 'intro' (string) and 'choices' (list of strings)."
    )
    response = llm.invoke({"messages": [{"role": "system", "content": prompt}]}).content
    data = json.loads(response)

    state["story"] = data["intro"]
    state["options"] = data["choices"]

    # Пауза для выбора игрока
    return interrupt(state, {
        "type": "choice",
        "question": f"{state['story']}\nВыберите действие:",
        "options": state["options"],
    })


def finish_story(state: State) -> State:
    """Генерирует окончание истории на основе выбранного варианта."""
    prompt = (
        f"Write an ending for the story based on the choice '{state['choice']}'. "
        "Keep it concise and satisfying."
    )
    response = llm.invoke({"messages": [{"role": "system", "content": prompt}]}).content
    state["ending"] = response
    return state


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

builder = StateGraph(State)
builder.add_node("generate", generate_story)
builder.add_node("finish", finish_story)

builder.set_entry_point("generate")
builder.add_edge(START, "generate")
builder.add_conditional_edges(
    "generate",
    lambda x: True if x.get("choice") else None,
    {True: "finish"},
)
builder.add_edge("finish", START)  # цикл можно завершить здесь

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


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

def main() -> None:
    thread_id = "interactive_story"

    init_state: State = {"story": "", "options": [], "choice": None, "ending": None}
    config = {"configurable": {"thread_id": thread_id}}

    # Запускаем поток генерации
    stream = graph.stream(Command(resume=init_state), config=config)

    for chunk in stream:
        if "__interrupt__" in chunk:
            payload = chunk["__interrupt__"][0].value
            answer = questionary.select(
                payload["question"],
                choices=payload["options"],
            ).ask()
            if answer is None:
                raise RuntimeError("Пользователь отменил ввод.")
            # Возобновляем граф с выбранным вариантом
            stream = graph.stream(Command(resume={"choice": answer}), config=config)
            continue

        # Выводим сообщения от LLM
        if "story" in chunk and chunk["story"]:
            console.print(f"[bold cyan]История:[/]\n{chunk['story']}")
        if "ending" in chunk and chunk["ending"]:
            console.print(f"\n[bold green]Концовка:[/]\n{chunk['ending']}")

    # После завершения выводим итоговую историю
    final_state = graph.get_state(config=config)
    full_story = f"{final_state['story']}\n\n{final_state['ending']}"
    console.print("\n[bold magenta]Итоговая история:[/]\n" + full_story)


if __name__ == "__main__":
    main()