import os
import uuid
import asyncio
import questionary
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage
from langchain.tools import tool
from langgraph.graph import StateGraph, START, END, Command
from langgraph.graph.message import add_messages
from langgraph.checkpoint.memory import InMemorySaver
from deepagents import create_deep_agent

# LLM configuration – 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.0,
)

# ---------- Graph State ----------
class GameState(TypedDict):
    messages: Annotated[list, add_messages]
    theme: str
    intro: str
    options: list[str]
    choice: str
    ending: str

# ---------- Graph Nodes ----------
async def generate_scene(state: GameState) -> GameState:
    theme = state["theme"]
    prompt = (
        f"Theme: {theme}\n"
        "Generate a short introduction (2-3 sentences) followed by exactly three numbered options for the hero. "
        "Output format: first line is the intro, next three lines are options numbered 1) 2) 3)."
    )
    response = await llm.ainvoke(HumanMessage(content=prompt))
    text = response.content.strip()
    lines = text.splitlines()
    intro = lines[0].strip()
    options = [line.strip() for line in lines[1:4]]
    return {**state, "intro": intro, "options": options}

async def interrupt_choice(state: GameState) -> GameState:
    payload = {
        "type": "choice",
        "question": f"{state['intro']}\nWhat do you do?",
        "options": state["options"],
    }
    # Pause execution until user responds
    await interrupt(payload)
    # After resume, the payload will contain 'choice'
    return {**state, "choice": payload["choice"]}

async def generate_ending(state: GameState) -> GameState:
    prompt = (
        f"Intro: {state['intro']}\n"
        f"Choice: {state['choice']}\n"
        "Write a short ending (2-3 sentences) for this story."
    )
    response = await llm.ainvoke(HumanMessage(content=prompt))
    ending = response.content.strip()
    return {**state, "ending": ending}

# ---------- Build Graph ----------
graph = StateGraph(GameState)
graph.add_node("scene", generate_scene)
graph.add_node("choice", interrupt_choice)
graph.add_node("ending", generate_ending)

graph.set_entry_point("scene")
graph.add_edge("scene", "choice")
graph.add_edge("choice", "ending")
graph.add_edge("ending", END)

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

# ---------- Game Runner ----------
async def run_game(theme: str) -> dict:
    thread_id = str(uuid.uuid4())
    config = {"configurable": {"thread_id": thread_id}}
    state: GameState = {"messages": [], "theme": theme, "intro": "", "options": [], "choice": "", "ending": ""}
    stream = graph.stream(state, config)
    last_state = state
    async for chunk in stream:
        # Handle interrupt
        if "__interrupt__" in chunk:
            interrupt_payload = chunk["__interrupt__"][0].value
            # Show options to user
            answer = questionary.select(
                interrupt_payload["question"],
                choices=interrupt_payload["options"],
            ).ask()
            # Resume with user's choice
            resume_payload = {**interrupt_payload, "choice": answer}
            stream = graph.stream(Command(resume=resume_payload), config)
            continue
        # Capture state when available
        if "state" in chunk:
            last_state = chunk["state"]
    return last_state

# ---------- DeepAgent Tool ----------
@tool
def play_game(theme: str) -> str:
    """Play a choose-your-own-adventure game with the given theme."""
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    final_state = loop.run_until_complete(run_game(theme))
    loop.close()
    intro = final_state["intro"]
    ending = final_state["ending"]
    return f"\n{intro}\n\n{ending}\n"

# ---------- Create DeepAgent ----------
agent = create_deep_agent(
    model=llm,
    tools=[play_game],
    backend=None,
    system_prompt="You are a game master that can play choose-your-own-adventure games.",
)

# ---------- CLI ----------
async def main():
    theme = input("Enter a theme for the adventure: ")
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=f"Play a game with theme: {theme}")]},
        {"configurable": {"thread_id": "cli-session"}},
    )
    # The tool returns the full story
    print(result["messages"][-1].content)

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