"""Интерактивная история: LLM + LangGraph interrupt + questionary."""
from __future__ import annotations

import os
import re
import uuid
from typing import TypedDict

import questionary
from dotenv import load_dotenv

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

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

LLM_MODEL = os.getenv("OPENROUTER_MODEL", "openai/gpt-oss-20b:free")
LLM_BASE_URL = os.getenv("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1")


class StoryState(TypedDict, total=False):
    theme: str
    setup: str
    scene_text: str
    options: list[str]
    user_choice: str
    ending: str


def build_llm() -> ChatOpenAI:
    return ChatOpenAI(
        model=LLM_MODEL,
        base_url=LLM_BASE_URL,
        api_key=os.getenv("OPENAI_API_KEY", "fake"),
        temperature=0.8,
    )


def _parse_scene_response(text: str) -> tuple[str, list[str]]:
    """Разбирает ответ LLM на завязку и 3 варианта."""
    options: list[str] = []
    setup = text.strip()

    variants_match = re.search(
        r"ВАРИАНТЫ:\s*(.*)",
        text,
        flags=re.IGNORECASE | re.DOTALL,
    )
    if variants_match:
        setup = text[: variants_match.start()].strip()
        block = variants_match.group(1)
        for line in block.splitlines():
            line = line.strip()
            if not line:
                continue
            cleaned = re.sub(r"^\d+[\).\]]\s*", "", line).strip()
            if cleaned:
                options.append(cleaned)

    if len(options) < 3:
        numbered = re.findall(r"^\s*\d+[\).\]]\s*(.+)$", text, flags=re.MULTILINE)
        options = [o.strip() for o in numbered if o.strip()]

    if len(options) < 3:
        parts = [p.strip() for p in re.split(r"[,;]", text) if p.strip()]
        if len(parts) >= 4:
            setup = parts[0]
            options = parts[1:4]

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

    return setup, options[:3]


def generate_scene(state: StoryState) -> dict:
    theme = state.get("theme", THEME_DEFAULT)
    llm = build_llm()
    prompt = (
        f"Тема: {theme}. Придумай короткую завязку (2–3 предложения) "
        "и ровно 3 варианта поступка героя.\n"
        "Формат ответа строго:\n"
        "ЗАВЯЗКА:\n"
        "<текст завязки>\n"
        "ВАРИАНТЫ:\n"
        "1) <вариант 1>\n"
        "2) <вариант 2>\n"
        "3) <вариант 3>"
    )
    response = llm.invoke(
        [
            SystemMessage(content="Ты автор интерактивных текстовых историй."),
            HumanMessage(content=prompt),
        ]
    )
    content = str(response.content)
    setup, options = _parse_scene_response(content)
    return {
        "theme": theme,
        "setup": setup,
        "scene_text": setup,
        "options": options,
    }


def choice_and_ending(state: StoryState) -> dict:
    setup = state.get("setup", "")
    options = state.get("options", [])

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

    if isinstance(resumed, dict):
        user_choice = (
            resumed.get("user_answer")
            or resumed.get("answer")
            or resumed.get("choice")
            or ""
        )
    else:
        user_choice = str(resumed)

    llm = build_llm()
    ending_prompt = (
        f"Завязка: {setup}\n"
        f"Выбор пользователя: {user_choice}\n"
        "Допиши короткую концовку (2–3 предложения). Только текст концовки."
    )
    ending_response = llm.invoke(
        [
            SystemMessage(content="Ты автор интерактивных историй."),
            HumanMessage(content=ending_prompt),
        ]
    )
    ending = str(ending_response.content).strip()

    return {
        "user_choice": user_choice,
        "ending": ending,
    }


def build_graph():
    builder = StateGraph(StoryState)
    builder.add_node("generate_scene", generate_scene)
    builder.add_node("choice_and_ending", choice_and_ending)
    builder.add_edge(START, "generate_scene")
    builder.add_edge("generate_scene", "choice_and_ending")
    memory = InMemorySaver()
    return builder.compile(checkpointer=memory)


def _handle_interrupt(interrupts: tuple) -> dict:
    first = interrupts[0]
    payload = first.value if hasattr(first, "value") else first
    if not isinstance(payload, dict):
        payload = {"type": "choice", "question": str(payload), "options": []}

    print(f"\n[LLM] {payload.get('question', '')}\n")

    options = payload.get("options") or []
    if options:
        choice = questionary.select(
            "Выберите действие:",
            choices=options,
        ).ask()
    else:
        choice = questionary.text("Ваш выбор:").ask()

    payload = dict(payload)
    payload["user_answer"] = choice or ""
    return payload


def run_story(theme: str = THEME_DEFAULT) -> StoryState:
    graph = build_graph()
    config = {"configurable": {"thread_id": str(uuid.uuid4())}}
    inputs: StoryState | Command = {"theme": theme}

    final_state: StoryState = {"theme": theme}

    while True:
        interrupted = False
        for chunk in graph.stream(inputs, config=config, stream_mode="updates"):
            if "__interrupt__" in chunk:
                payload = _handle_interrupt(chunk["__interrupt__"])
                inputs = Command(resume=payload)
                interrupted = True
                break

            for node_name, update in chunk.items():
                if node_name == "__interrupt__":
                    continue
                if isinstance(update, dict):
                    final_state.update(update)
                    if node_name == "generate_scene" and update.get("setup"):
                        print(f"\nТема: {theme}")
                        print(f"\n[LLM] {update['setup']}\n")
                    if node_name == "choice_and_ending" and update.get("ending"):
                        print(f"\n[LLM] {update['ending']}\n")

        if not interrupted:
            break

    print("--- Итог ---")
    print(f"Завязка: {final_state.get('setup', '')}")
    print(f"Выбор: {final_state.get('user_choice', '')}")
    print(f"Концовка: {final_state.get('ending', '')}")
    return final_state


if __name__ == "__main__":
    import sys

    story_theme = sys.argv[1] if len(sys.argv) > 1 else THEME_DEFAULT
    run_story(story_theme)