import os
import asyncio
import questionary
from typing import TypedDict, Annotated
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, START, END, interrupt, Command
from langgraph.graph.message import add_messages
from langgraph.checkpoint.memory import InMemorySaver
from langchain_core.messages import HumanMessage, AIMessage
from deepagents import create_deep_agent
from deepagents.backends import CompositeBackend, LocalShellBackend, FilesystemBackend
from deepagents.tools import tool

# ---------- LLM ----------
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):
    theme: str
    scene: Annotated[str, add_messages]
    options: list[str]
    choice: str
    ending: Annotated[str, add_messages]

# ---------- Graph nodes ----------
async def generate_scene(state: GameState) -> GameState:
    theme = state["theme"]
    prompt = (
        f"Тема: {theme}.\n"
        "Придумай короткую завязку (2–3 предложения) и ровно 3 варианта поступка героя.\n"
        "Ответь в формате: сначала текст завязки, затем строка с вариантами через запятую."
    )
    response = await llm.ainvoke(HumanMessage(content=prompt))
    text = response.content
    # Разделяем завязку и варианты
    parts = text.split("\n")
    scene_text = parts[0].strip()
    options_line = parts[1] if len(parts) > 1 else ""
    options = [opt.strip() for opt in options_line.split(",") if opt.strip()]
    state["scene"] = [AIMessage(content=scene_text)]
    state["options"] = options
    # Прерываем для выбора
    interrupt_payload = {
        "type": "choice",
        "question": f"{scene_text}\n\nЧто делаем?",
        "options": options,
    }
    interrupt(interrupt_payload)
    return state

async def finish_story(state: GameState) -> GameState:
    # state now contains "choice"
    scene_text = state["scene"][0].content
    choice = state["choice"]
    prompt = (
        f"Завязка: {scene_text}\n"
        f"Выбор пользователя: {choice}\n"
        "Допиши короткую концовку (2–3 предложения)."
    )
    response = await llm.ainvoke(HumanMessage(content=prompt))
    ending_text = response.content
    state["ending"] = [AIMessage(content=ending_text)]
    return state

# ---------- Build graph ----------
builder = StateGraph(GameState)
builder.add_node("generate", generate_scene)
builder.add_node("finish", finish_story)
builder.set_entry_point("generate")
builder.add_edge("generate", "finish")
builder.add_edge("finish", END)

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

# ---------- Tool that runs the game ----------
@tool
async def play_game(theme: str) -> str:
    """Play a choose‑your‑adventure game with the given theme."""
    thread_id = f"game-{theme.replace(' ', '-')}-1"
    config = {"configurable": {"thread_id": thread_id}}
    # Initial state
    state: GameState = {"theme": theme, "scene": [], "options": [], "choice": "", "ending": []}
    # Start streaming
    stream = graph.stream(state, config)
    async for chunk in stream:
        if "__interrupt__" in chunk:
            # Extract payload
            payload = chunk["__interrupt__"][0].value
            # Show question and options
            answer = questionary.select(
                payload["question"],
                choices=payload["options"],
            ).ask()
            # Resume with answer
            payload["choice"] = answer
            resume_cmd = Command(resume=payload)
            stream = graph.stream(resume_cmd, config)
            continue
        # Collect messages
        if "scene" in chunk:
            state["scene"] = chunk["scene"]
        if "ending" in chunk:
            state["ending"] = chunk["ending"]
    # Build final story
    story = "\n\n".join([m.content for m in state["scene"] + state["ending"]])
    return story

# ---------- DeepAgent setup ----------
backend = CompositeBackend([
    LocalShellBackend(workspace_dir="./workspace"),
    FilesystemBackend(),
])

agent = create_deep_agent(
    model=llm,
    tools=[play_game],
    backend=backend,
    system_prompt="You are a helpful assistant that can play choose‑your‑adventure games.",
)

async def main():
    theme = questionary.text("Enter a theme for the adventure:").ask()
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=f"play_game {theme}")]},
        {"configurable": {"thread_id": "session-1"}},
    )
    print("\n\n--- Final Story ---\n")
    print(result["messages"][-1].content)

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