import os
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, START
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.prebuilt import interrupt_and_resume
from typing import TypedDict, List
import questionary

class GameState(TypedDict):
    intro: str | None
    options: List[str] | None
    choice: str | None
    ending: str | None

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7)

async def generate_intro_and_options(state: GameState):
    prompt = (
        "You are a creative storyteller. Generate an exciting intro for a choose-your-own-adventure game and 3 distinct options the player can choose from. Return JSON with keys 'intro' and 'options'."
    )
    res = await llm.ainvoke(prompt)
    import json
    data = json.loads(res.content)
    state["intro"] = data.get("intro")
    state["options"] = data.get("options")
    return interrupt_and_resume(state, {
        "type": "question",
        "text": data.get("intro"),
        "choices": data.get("options"),
    })

async def generate_ending(state: GameState):
    prompt = (
        f"Player chose: {state['choice']}. Write a short ending that follows the intro and incorporates the choice."
    )
    res = await llm.ainvoke(prompt)
    state["ending"] = res.content
    return state

builder = StateGraph(GameState)
builder.add_node("intro", generate_intro_and_options)
builder.add_node("ending", generate_ending)
builder.set_entry_point("intro")
builder.add_edge(START, "intro")
builder.add_conditional_edges(
    "intro",
    lambda x: "choice" if isinstance(x.get("__interrupt__"), dict) else None,
)
builder.add_edge("choice", "ending")
graph = builder.compile(checkpointer=InMemorySaver())

async def main():
    state = await graph.astart()
    while True:
        if "__interrupt__" in state:
            intr = state["__interrupt__"]
            answer = questionary.select(intr["text"], choices=intr["choices"]).ask()
            state["choice"] = answer
            state = await graph.arun(state)
        else:
            break
    print("\n--- Ending ---")
    print(state.get("ending"))

if __name__ == "__main__":
    import asyncio
    asyncio.run(main())