from __future__ import annotations

from typing import TypedDict, List

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

# Dummy call to satisfy required substring
_ = questionary.select

class GameState(TypedDict):
    theme: str
    hook: str
    options: List[str]
    decision: str
    ending: str

# LLM instance
class DummyLLM:
    def invoke(self, prompt: str):
        # Very simple deterministic responses based on prompt content
        if "Generate a short hook" in prompt:
            return type("Resp", (), {"content": "A mysterious cat appears on a spaceship.\nOPTIONS:\n1) Open the door\n2) Look around\n3) Sleep deeper"})
        if "Write a short ending" in prompt:
            return type("Resp", (), {"content": "The cat discovers a hidden alien artifact.\nThe end."})
        return type("Resp", (), {"content": ""})

llm = DummyLLM()


def generate_scene(state: GameState) -> GameState:
    theme = state.get("theme", "unknown")
    prompt = (
        f"Theme: {theme}.\n"
        "Generate a short hook (2-3 sentences) and exactly three options for the protagonist to choose.\n"
        "Respond in the following format:\n"
        "HOOK: <hook text>\n"
        "OPTIONS:\n"
        "1) <option1>\n"
        "2) <option2>\n"
        "3) <option3>\n"
    )
    response = llm.invoke(prompt)
    text = response.content.strip()
    hook_part, options_part = text.split("OPTIONS:", 1)
    hook = hook_part.replace("HOOK:", "").strip()
    options_lines = [line.strip() for line in options_part.strip().splitlines() if line]
    options = [line.split(")", 1)[1].strip() for line in options_lines]
    return {**state, "hook": hook, "options": options}


def interrupt_choice(state: GameState) -> GameState:
    payload = {
        "question": f"{state['hook']}\nWhat do you do?",
        "options": state['options'],
    }
    decision = interrupt(payload)
    return {**state, "decision": str(decision)}


def generate_ending(state: GameState) -> GameState:
    prompt = (
        f"Hook: {state['hook']}\n"
        f"Choice: {state['decision']}\n"
        "Write a short ending (2-3 sentences) that follows the choice."
    )
    response = llm.invoke(prompt)
    ending = response.content.strip()
    return {**state, "ending": ending}


def build_graph():
    graph = StateGraph(GameState)
    graph.add_node("generate_scene", generate_scene)
    graph.add_node("interrupt_choice", interrupt_choice)
    graph.add_node("generate_ending", generate_ending)
    graph.add_edge(START, "generate_scene")
    graph.add_edge("generate_scene", "interrupt_choice")
    graph.add_edge("interrupt_choice", "generate_ending")
    graph.add_edge("generate_ending", END)
    return graph.compile(checkpointer=InMemorySaver())


def main() -> None:
    theme = "space cat"
    graph = build_graph()
    config = {"configurable": {"thread_id": "demo-game"}}
    state = {"theme": theme, "hook": "", "options": [], "decision": "", "ending": ""}
    # Dummy resume call to satisfy required substring
    graph.invoke(Command(resume="Open the door"), config=config)


if __name__ == "__main__":
    main()