from langchain_openai import ChatOpenAI
from pydantic import SecretStr
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.types import interrupt, Command
import questionary

# LLM configuration with placeholders
llm = ChatOpenAI(
    model="openai/gpt-oss-20b",
    base_url='https://platform.brojs.ru/jrnl-bh/api/inference/v1',
    api_key=SecretStr("jrnl_30283ab953615cbb6846ff9940a1eedce0b76d7b2f59a2394f29e74643e6a90d"),
    temperature=0.7,
)

# State definition
class GameState(dict):
    topic: str | None = None
    intro: str | None = None
    options: list[str] | None = None
    choice: str | None = None
    ending: str | None = None

# Node to generate intro and options
def generate_intro(state: GameState) -> GameState:
    topic = state.get("topic", "неизвестная тема")
    prompt = f"""Тема: {topic}. Придумай короткую завязку (2–3 предложения) и ровно 3 варианта действий героя. Ответь в формате:
Завязка:
<текст>
Варианты:
1) <вариант1>
2) <вариант2>
3) <вариант3>"""
    response = llm.invoke(prompt).content
    # Parse response
    parts = response.split("Варианты:")
    intro_part = parts[0].replace("Завязка:", "").strip()
    options_part = parts[1] if len(parts) > 1 else ""
    options = [opt.strip() for opt in options_part.splitlines() if opt.strip()]
    state["intro"] = intro_part
    state["options"] = options
    # Interrupt to ask user
    return interrupt(
        {
            "type": "choice",
            "question": f"{intro_part}\n\nЧто делаем?",
            "options": options,
        }
    )

# Node to handle resume after choice and generate ending
def generate_ending(state: GameState) -> GameState:
    intro = state.get("intro", "")
    choice = state.get("choice", "")
    prompt = f"""Завязка: {intro}
Выбор пользователя: {choice}
Допиши короткую концовку (2–3 предложения)."""
    ending = llm.invoke(prompt).content.strip()
    state["ending"] = ending
    return state

# Build graph
graph_builder = StateGraph(GameState)
graph_builder.add_node("generate_intro", generate_intro)
graph_builder.add_node("generate_ending", generate_ending)
graph_builder.set_entry_point("generate_intro")
graph_builder.add_edge(START, "generate_intro")
graph_builder.add_edge("generate_intro", "generate_ending")  # continue after interrupt
graph_builder.add_edge("generate_ending", END)

# Compile with checkpoint
checkpoint = InMemorySaver()
app = graph_builder.compile(checkpointer=checkpoint)

def run_game(topic: str):
    state = GameState({"topic": topic})
    config = {"configurable": {"thread_id": "game_thread"}}
    # First stream until interrupt
    for chunk in app.stream(state, config):
        if "__interrupt__" in chunk:
            payload = chunk["__interrupt__"][0].value  # dict from interrupt()
            answer = questionary.select(
                payload["question"],
                choices=payload["options"]
            ).ask()
            # Resume with user's choice
            resume_payload = {"choice": answer}
            for res_chunk in app.stream(Command(resume=resume_payload), config):
                if "ending" in res_chunk:
                    print("\n[LLM] Концовка:", res_chunk["ending"])
    # Final state
    final_state = checkpoint.get("game_thread")
    print("\nИтоговое состояние:")
    print("Тема:", final_state.get("topic"))
    print("Завязка:", final_state.get("intro"))
    print("Выбор:", final_state.get("choice"))
    print("Концовка:", final_state.get("ending"))

if __name__ == "__main__":
    topic_input = input("Введите тему истории: ").strip()
    if not topic_input:
        topic_input = "космический кот"
    run_game(topic_input)