import os
import questionary
from typing import TypedDict, List, Optional

from langgraph.graph import StateGraph, START, END
from langgraph.types import interrupt, Command
from langgraph.checkpoint.memory import InMemorySaver

from langchain_openai import ChatOpenAI

# LLM setup – OpenRouter
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://openrouter.ai/api/v1",
    api_key=os.getenv("OPENAI_API_KEY"),
    temperature=0.7,
)

# State definition
class GameState(TypedDict):
    theme: str
    scene: Optional[str]
    options: Optional[List[str]]
    choice: Optional[str]
    ending: Optional[str]

# Node: generate scene and options
async def generate_scene(state: GameState) -> GameState:
    prompt = (
        f"Тема: {state['theme']}\n"
        "Придумай короткую завязку (2–3 предложения) и ровно 3 варианта поступка героя.\n"
        "Ответь в формате: сначала текст завязки, затем каждая строка с вариантом,\n"
        "по одному на строку, без нумерации."
    )
    response = await llm.ainvoke([{"role": "user", "content": prompt}])
    text = response.content.strip()
    parts = text.split("\n")
    # first non-empty line is scene
    scene = None
    options: List[str] = []
    for line in parts:
        line = line.strip()
        if not line:
            continue
        if scene is None:
            scene = line
        else:
            options.append(line)
    state["scene"] = scene
    state["options"] = options
    # interrupt for choice
    payload = {
        "type": "choice",
        "question": f"{scene}\n\nЧто делаем?",
        "options": options,
    }
    return interrupt(payload)

# Node: finish story after choice
async def finish_story(state: GameState) -> GameState:
    # state now contains 'choice'
    prompt = (
        f"Завязка: {state['scene']}\n"
        f"Выбор пользователя: {state['choice']}\n"
        "Допиши короткую концовку (2–3 предложения)."
    )
    response = await llm.ainvoke([{"role": "user", "content": prompt}])
    ending = response.content.strip()
    state["ending"] = ending
    return state

# Build graph
builder = StateGraph(GameState)
builder.add_node("generate", generate_scene)
builder.add_node("finish", finish_story)
builder.add_edge(START, "generate")
builder.add_edge("generate", "finish")
builder.add_edge("finish", END)

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

# Runner
async def run_game(theme: str):
    state: GameState = {
        "theme": theme,
        "scene": None,
        "options": None,
        "choice": None,
        "ending": None,
    }
    config = {"configurable": {"thread_id": "game1"}}
    # start stream
    stream = graph.stream(state, config)
    async for chunk in stream:
        if "__interrupt__" in chunk:
            # handle interrupt
            interrupt_payload = chunk["__interrupt__"][0].value
            # show question and options
            answer = questionary.select(
                interrupt_payload["question"],
                choices=interrupt_payload["options"],
            ).ask()
            # resume with answer
            resume_payload = {**interrupt_payload, "choice": answer}
            stream = graph.stream(Command(resume=resume_payload), config)
            continue
        # print normal messages
        if "messages" in chunk:
            for msg in chunk["messages"]:
                if msg.role == "assistant":
                    print(msg.content)
    # after completion, print final state
    print("\nИтоговое состояние:")
    print(state)

if __name__ == "__main__":
    import asyncio
    theme = questionary.text("Введите тему игры:").ask()
    asyncio.run(run_game(theme))