"""
Interactive choose-your-own-adventure game using LangGraph and OpenAI.

Run with:
    python game.py

Make sure you have an OpenAI API key set in the environment variable OPENAI_API_KEY.
"""

import os
import re
from typing import TypedDict, List, Dict, Any

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

# ---------------------------------------------------------------------------
# 1. Define the graph state
# ---------------------------------------------------------------------------
class GameState(TypedDict):
    theme: str
    story: str
    options: List[str]
    choice: str
    ending: str

# ---------------------------------------------------------------------------
# 2. Helper functions
# ---------------------------------------------------------------------------

# Simple parser that expects the LLM to return:
#   <story>
#   1) <option1>
#   2) <option2>
#   3) <option3>
# The story may span multiple lines.

def parse_story_and_options(text: str) -> tuple[str, List[str]]:
    """Parse the LLM output into a story and a list of options.

    The function is tolerant to small formatting variations.
    """
    # Split into lines
    lines = text.strip().splitlines()
    # Find the first line that starts with a digit and a closing parenthesis or dot
    option_pattern = re.compile(r"^\s*\d+[\).]\s*(.*)")
    story_lines = []
    options: List[str] = []
    for line in lines:
        m = option_pattern.match(line)
        if m:
            options.append(m.group(1).strip())
        else:
            story_lines.append(line)
    story = "\n".join(story_lines).strip()
    return story, options

# ---------------------------------------------------------------------------
# 3. Graph nodes
# ---------------------------------------------------------------------------

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

def generate_scene_and_interrupt(state: GameState) -> GameState:
    """Generate the opening scene and options, then interrupt for user choice.

    The node returns a new state with story and options populated, and then
    calls interrupt(). After the user responds, the graph resumes in this
    same node with the updated state containing the choice.
    """
    theme = state["theme"]

    # If we have already received a choice, generate the ending and finish.
    if state.get("choice"):
        # Generate ending
        prompt = (
            f"Theme: {theme}\n"
            f"Story: {state['story']}\n"
            f"User choice: {state['choice']}\n"
            "Write a short ending (2–3 sentences) that follows from the choice."
        )
        ending = llm.invoke(prompt).content
        state["ending"] = ending.strip()
        return state

    # No choice yet: generate scene and options
    prompt = (
        f"Theme: {theme}\n"
        "Create a short opening scene (2–3 sentences) and exactly three numbered options for the hero to choose.\n"
        "Respond in the following format:\n"
        "<story>\n"
        "1) <option1>\n"
        "2) <option2>\n"
        "3) <option3>\n"
    )
    raw = llm.invoke(prompt).content
    story, options = parse_story_and_options(raw)
    state["story"] = story
    state["options"] = options

    # Prepare interrupt payload
    interrupt_payload = {
        "type": "choice",
        "question": story,
        "options": options,
    }
    # Interrupt the graph; the returned state will contain the user response
    interrupt(interrupt_payload)
    # After interrupt, the graph will resume in this same node with the updated state.
    return state

# ---------------------------------------------------------------------------
# 4. Build the graph
# ---------------------------------------------------------------------------

graph_builder = StateGraph(GameState)
graph_builder.add_node("scene", generate_scene_and_interrupt)
# All paths go through the same node; we finish when ending is set.
graph_builder.set_entry_point("scene")
# The node will keep looping until an ending is produced.
# We use a simple condition: if ending exists, we finish.

def is_finished(state: GameState) -> bool:
    return bool(state.get("ending"))

graph_builder.add_conditional_edges("scene", lambda s: "finished" if is_finished(s) else "scene")
graph_builder.add_edge("finished", "finished")  # terminal

# Compile the graph with a checkpoint
checkpoint = InMemorySaver()
graph = graph_builder.compile(checkpointer=checkpoint)

# ---------------------------------------------------------------------------
# 5. Main loop handling interrupts
# ---------------------------------------------------------------------------

def run_game(theme: str):
    # Initial state
    state: GameState = {
        "theme": theme,
        "story": "",
        "options": [],
        "choice": "",
        "ending": "",
    }

    thread_id = os.urandom(8).hex()
    config = {"configurable": {"thread_id": thread_id}}

    # Start streaming
    stream = graph.stream(Command(state), config)
    try:
        for chunk in stream:
            # The chunk may contain an interrupt
            if "__interrupt__" in chunk:
                interrupt_payload = chunk["__interrupt__"][0].value
                # Show question and options to user
                answer = questionary.select(
                    interrupt_payload["question"],
                    choices=interrupt_payload["options"],
                ).ask()
                # Add the answer to the payload
                interrupt_payload["choice"] = answer
                # Resume the graph with the updated payload
                stream = graph.stream(Command(resume=interrupt_payload), config)
                continue
            # Normal output: print to console
            if "story" in chunk:
                print(chunk["story"], end="\n\n")
            if "ending" in chunk:
                print("\n[LLM] " + chunk["ending"], end="\n\n")
    except KeyboardInterrupt:
        print("\nGame interrupted by user.")

    # Final state
    final_state = graph.get_state(config)
    print("--- Final state ---")
    print("Theme:", final_state["theme"])
    print("Story:", final_state["story"])
    print("Choice:", final_state["choice"])
    print("Ending:", final_state["ending"])

if __name__ == "__main__":
    print("Welcome to the LLM choose‑your‑own‑adventure game!")
    theme = questionary.text("Enter a theme for the story:").ask()
    if not theme:
        theme = "A mysterious space cat"
    run_game(theme)