"""Текстовая игра «выбери свою историю»: LLM + interrupt (LangGraph)."""
from __future__ import annotations

import os
import re
import sys
from typing import TypedDict

import questionary
from dotenv import load_dotenv
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.constants import START
from langgraph.graph import StateGraph
from langgraph.types import Command, interrupt

load_dotenv()

llm = ChatOpenAI(
    model=os.getenv("OPENAI_MODEL", "openai/gpt-oss-20b:free"),
    base_url=os.getenv("OPENAI_BASE_URL", "https://openrouter.ai/api/v1"),
    api_key=os.getenv("OPENAI_API_KEY", "fake"),
    temperature=0.8,
)

DEFAULT_THEME = "космический кот"


class StoryState(TypedDict):
    theme: str
    scene_text: str
    choices: list[str]
    user_choice: str
    ending: str


def _call_llm(system: str, user: str) -> str:
    response = llm.invoke([SystemMessage(content=system), HumanMessage(content=user)])
    return str(response.content).strip()


def _parse_setup(raw: str) -> tuple[str, list[str]]:
    """Разбор ответа LLM: завязка + три варианта."""
    setup = raw
    choices: list[str] = []

    if "ВАРИАНТЫ:" in raw:
        setup, variants_block = raw.split("ВАРИАНТЫ:", maxsplit=1)
        setup = setup.replace("ЗАВЯЗКА:", "").strip()
        for line in variants_block.strip().splitlines():
            line = re.sub(r"^\d+[\).\]]\s*", "", line.strip())
            if line:
                choices.append(line)
    else:
        numbered = re.findall(r"(?:^|\n)\s*\d+[\).\]]\s*(.+)", raw)
        if numbered:
            choices = [c.strip() for c in numbered[:3]]
            setup = re.split(r"\n\s*1[\).\]]", raw, maxsplit=1)[0].strip()

    if len(choices) < 3:
        for line in raw.splitlines():
            if line.strip().startswith(("-", "•")):
                choices.append(line.strip().lstrip("-• ").strip())
        choices = choices[:3]

    while len(choices) < 3:
        choices.append(f"Вариант {len(choices) + 1}")

    return setup.strip(), choices[:3]


def story_node(state: StoryState) -> StoryState:
    """Генерация завязки → прерывание (выбор) → концовка от LLM."""
    theme = state.get("theme") or DEFAULT_THEME

    setup_raw = _call_llm(
        system=(
            "Ты автор интерактивных текстовых игр. Отвечай строго в формате:\n"
            "ЗАВЯЗКА:\n<2-3 предложения>\n\n"
            "ВАРИАНТЫ:\n1) <вариант 1>\n2) <вариант 2>\n3) <вариант 3>"
        ),
        user=f"Тема: {theme}. Придумай короткую завязку и ровно 3 варианта поступка героя.",
    )
    scene_text, choices = _parse_setup(setup_raw)

    print(f"\n[LLM] {scene_text}\n")

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

    user_choice = payload.get("answer", choices[0])
    print(f"\n> {user_choice}\n")

    ending = _call_llm(
        system="Ты автор интерактивных историй. Допиши короткую концовку (2–3 предложения).",
        user=(
            f"Тема: {theme}\n"
            f"Завязка: {scene_text}\n"
            f"Выбор пользователя: {user_choice}\n"
            "Допиши короткую концовку с учётом выбора."
        ),
    )

    print(f"[LLM] {ending}\n")

    return {
        "theme": theme,
        "scene_text": scene_text,
        "choices": choices,
        "user_choice": user_choice,
        "ending": ending,
    }


def _ask_user(payload: dict) -> dict:
    question = payload.get("question", "Что делаем?")
    options = payload.get("allow_responds") or ["Вариант 1", "Вариант 2", "Вариант 3"]
    choice = questionary.select(question, choices=options).ask()
    if choice is None:
        choice = options[0]
    updated = dict(payload)
    updated["answer"] = choice
    return updated


def play_story(theme: str) -> StoryState:
    builder = StateGraph(StoryState)
    builder.add_node("story", story_node)
    builder.add_edge(START, "story")
    graph = builder.compile(checkpointer=InMemorySaver())

    config = {"configurable": {"thread_id": f"story-{theme}"}}
    initial: StoryState = {
        "theme": theme,
        "scene_text": "",
        "choices": [],
        "user_choice": "",
        "ending": "",
    }

    final: StoryState = initial
    stream = graph.stream(initial, config)

    for chunk in stream:
        if "__interrupt__" in chunk:
            payload = dict(chunk["__interrupt__"][0].value)
            resumed = _ask_user(payload)
            resume_stream = graph.stream(Command(resume=resumed), config)
            for resume_chunk in resume_stream:
                if "story" in resume_chunk:
                    final = resume_chunk["story"]
        elif "story" in chunk:
            final = chunk["story"]

    return final


def main() -> None:
    theme = " ".join(sys.argv[1:]).strip() or DEFAULT_THEME
    print(f"Тема: {theme}\n")
    result = play_story(theme)

    print("=== Итоговая история ===")
    print(f"Тема: {result['theme']}")
    print(f"Завязка: {result['scene_text']}")
    print(f"Выбор: {result['user_choice']}")
    print(f"Концовка: {result['ending']}")


if __name__ == "__main__":
    main()