import os
import asyncio
from typing import TypedDict, List, Annotated

from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint import InMemorySaver
from langgraph.types import interrupt, Command
from langgraph.graph.message import add_messages

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langchain.tools import tool
from deepagents import create_deep_agent
from deepagents.backends import CompositeBackend, LocalShellBackend, FilesystemBackend

import questionary

# ---------- 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,
)

# ---------- Backend & Tools (required by deepagents) ----------
backend = CompositeBackend(
    [
        LocalShellBackend(workspace_dir="./workspace"),
        FilesystemBackend(),
    ]
)

@tool
def dummy_tool(query: str) -> str:
    """A placeholder tool required by the deep agent."""
    return f"tool result for {query}"

agent = create_deep_agent(
    model=llm,
    tools=[dummy_tool],
    backend=backend,
    system_prompt="You are a helpful storytelling assistant.",
)

# ---------- State ----------
class StoryState(TypedDict):
    messages: Annotated[List, add_messages]
    topic: str
    intro: str
    options: List[str]
    choice: str
    ending: str

# ---------- Node ----------
async def story_node(state: StoryState):
    # Phase 1: generate intro and options
    if "intro" not in state or not state["intro"]:
        prompt = (
            f"Topic: {state['topic']}. "
            "Create a short story beginning (2-3 sentences) and exactly three possible actions for the hero. "
            "Respond in the following format:\n"
            "Intro: <your text>\n"
            "Options:\n"
            "1) <first action>\n"
            "2) <second action>\n"
            "3) <third action>"
        )
        response = await agent.ainvoke(
            {"messages": [HumanMessage(content=prompt)]},
            {"configurable": {"thread_id": "agent-1"}},
        )
        text = response["messages"][-1].content

        # Simple parsing
        intro_part, options_part = text.split("Options:", 1)
        intro = intro_part.replace("Intro:", "").strip()
        raw_options = options_part.strip().splitlines()
        options = [line.split(")", 1)[1].strip() for line in raw_options if ")" in line]

        # Store and interrupt
        state["intro"] = intro
        state["options"] = options
        payload = {
            "type": "choice",
            "question": f"{intro}\n\nWhat does the hero do?",
            "options": options,
        }
        return interrupt(payload)

    # Phase 2: after user choice, generate ending
    if "choice" in state and state["choice"]:
        prompt = (
            f"Intro: {state['intro']}\n"
            f"User choice: {state['choice']}\n"
            "Write a short conclusion (2-3 sentences) that follows from this choice."
        )
        response = await agent.ainvoke(
            {"messages": [HumanMessage(content=prompt)]},
            {"configurable": {"thread_id": "agent-2"}},
        )
        ending = response["messages"][-1].content.strip()
        state["ending"] = ending
        return state

    # Should not reach here
    return state

# ---------- Graph ----------
graph = StateGraph(StoryState)
graph.add_node("story", story_node)
graph.add_edge(START, "story")
graph.add_edge("story", END)
graph.set_entry_point("story")
checkpointer = InMemorySaver()
graph = graph.compile(checkpointer=checkpointer)

# ---------- Main loop ----------
async def main():
    topic = questionary.text("Enter a story theme (e.g., 'space cat'):", default="space cat").ask()
    thread_id = "session-1"
    config = {"configurable": {"thread_id": thread_id}}

    # Initial empty state
    initial_state: StoryState = {
        "messages": [],
        "topic": topic,
        "intro": "",
        "options": [],
        "choice": "",
        "ending": "",
    }

    # First run - will pause at interrupt
    async for event in graph.stream(initial_state, config):
        if "__interrupt__" in event:
            payload = event["__interrupt__"][0].value
            question = payload["question"]
            choices = payload["options"]
            answer = questionary.select(question, choices=choices).ask()
            payload["choice"] = answer

            # Resume graph with the updated payload
            resume_cmd = Command(resume=payload)
            async for resume_event in graph.stream(resume_cmd, config):
                if "__interrupt__" in resume_event:
                    # No further interrupts expected
                    continue
                # Continue until END
            break

    # Retrieve final state
    final_state = await graph.aget_state(thread_id)
    print("\n--- Final Story ---")
    print(f"Theme: {final_state.state['topic']}\n")
    print(final_state.state["intro"])
    print(f"\nYou chose: {final_state.state['choice']}")
    print(f"\n{final_state.state['ending']}\n")

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