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

# LangGraph imports
from langgraph.graph import StateGraph, START
from langgraph.types import interrupt, Command
from langgraph.checkpoint.memory import InMemorySaver

# LangChain LLM
from langchain_openai import ChatOpenAI

# Console interaction
import questionary

# ---------- 1. Define state ---------------------------------
class GameState(TypedDict):
    theme: str
    intro: str
    options: List[str]
    choice: str | None
    ending: str

# ---------- 2. LLM client -----------------------------------
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://platform.brojs.ru/jrnl-bh/api/inference/v1",
    api_key=os.getenv("JOURNAL_MCP_PAT"),
    temperature=0.7,
)

# ---------- 3. Node to generate intro and options ---------
async def generate_intro(state: GameState) -> Dict:
    theme = state["theme"]
    prompt = (
        f"You are a creative storyteller.\n"
        f"Theme: {theme}.\n"
        "Generate a short opening paragraph (2-3 sentences).\n"
        "Then provide exactly three distinct actions the protagonist can take, one per line, prefixed with numbers 1) 2) 3).\n"
        "Return only the intro followed by the numbered list."
    )
    response = await llm.ainvoke([HumanMessage(content=prompt)])
    text = response.content
    # Split into intro and options
    parts = text.split("\n")
    intro_lines = []
    options: List[str] = []
    for line in parts:
        if line.strip().startswith(tuple(str(i)+")" for i in range(1,4)):
            options.append(line.strip())
        else:
            intro_lines.append(line)
    state["intro"] = " ".join(intro_lines).strip()
    state["options"] = options
    # Trigger interrupt to ask user
    payload = {
        "type": "choice",
        "question": f"{state['intro']}\n\nWhat do you do?",
        "choices": options,
    }
    return interrupt(payload)

# ---------- 4. Node to process choice and generate ending -----
async def generate_ending(state: GameState) -> Dict:
    # state now contains 'choice'
    theme = state["theme"]
    intro = state["intro"]
    choice = state.get("choice", "")
    prompt = (
        f"You are a creative storyteller.\n"
        f"Theme: {theme}.\n"
        f"Intro: {intro}.\n"
        f"The player chose: {choice}.\n"
        "Write a short ending (2-3 sentences) that follows from this choice."
    )
    response = await llm.ainvoke([HumanMessage(content=prompt)])
    state["ending"] = response.content.strip()
    return state

# ---------- 5. Build graph ---------------------------------
builder = StateGraph(GameState)
builder.add_node("intro", generate_intro)
builder.add_node("ending", generate_ending)
builder.set_entry_point("intro")
builder.add_edge(START, "intro")
builder.add_edge("intro", "ending")
# No edge from ending; graph ends after ending node
graph = builder.compile(checkpointer=InMemorySaver())

# ---------- 6. Run loop with interrupt handling ---------
async def run_game(theme: str):
    thread_id = f"thread-{theme.replace(' ', '_')}"
    config = {"configurable": {"thread_id": thread_id}}
    # Start stream
    async for event in graph.astream({"theme": theme}, config, stream_mode="messages"):
        if isinstance(event, tuple):
            msg, _meta = event
            print(msg.content)
        elif "__interrupt__" in event:
            interrupt_payload = event["__interrupt__"][0].value
            # Show question and choices using questionary
            answer = questionary.select(
                interrupt_payload["question"],
                choices=[c.split(') ',1)[1] if ') ' in c else c for c in interrupt_payload["choices"]]
            ).ask()
            # Map back to full choice string
            full_choice = next(c for c in interrupt_payload["choices"] if c.endswith(answer))
            # Resume with user answer
            resume_payload = {**interrupt_payload, "choice": full_choice}
            async for sub_event in graph.astream(Command(resume=resume_payload), config, stream_mode="messages"):
                if isinstance(sub_event, tuple):
                    msg, _meta = sub_event
                    print(msg.content)
    # After completion, print final state
    final_state = await graph.aget_state(config)
    print("\n--- Final Story ---")
    print(final_state["intro"])
    print(f"Choice: {final_state['choice']}")
    print(final_state["ending"])

if __name__ == "__main__":
    theme = questionary.text("Enter a theme for the adventure:").ask()
    asyncio.run(run_game(theme))