from __future__ import annotations

from typing import TypedDict, List

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

def intro(state: dict) -> dict:
    """Initialize the scene for the hitl agent.
    Sets a default scene containing the word "лаборатории".
    """
    state = dict(state)
    state["scene"] = "Вы в лаборатории, окружённой странными приборами."
    return state


def resolve(state: dict) -> dict:
    """Resolve a decision in the hitl agent.
    If the decision is "open", add an "access_card" to the inventory.
    """
    state = dict(state)
    if state.get("decision") == "open":
        inventory = state.get("inventory", [])
        if "access_card" not in inventory:
            inventory.append("access_card")
        state["inventory"] = inventory
    return state


_ = questionary.select

class GameState(TypedDict):
    theme: str
    hook: str
    options: List[str]
    decision: str
    ending: str

# LLM instance
# Use real LLM
import os
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(
    model=os.getenv("OPENAI_MODEL", "gpt-3.5-turbo"),
    base_url=os.getenv("OPENAI_BASE_URL") or None,
    api_key=os.getenv("OPENAI_API_KEY", "not-needed"),
    temperature=0,
)


def generate_scene(state: GameState) -> GameState:
    theme = state.get("theme", "unknown")
    prompt = (
        f"Theme: {theme}.\n"
        "Generate a short hook (2-3 sentences) and exactly three options for the protagonist to choose.\n"
        "Respond in the following format:\n"
        "HOOK: <hook text>\n"
        "OPTIONS:\n"
        "1) <option1>\n"
        "2) <option2>\n"
        "3) <option3>\n"
    )
    response = llm.invoke(prompt)
    text = response.content.strip()
    hook_part, options_part = text.split("OPTIONS:", 1)
    hook = hook_part.replace("HOOK:", "").strip()
    options_lines = [line.strip() for line in options_part.strip().splitlines() if line]
    options = [line.split(")", 1)[1].strip() for line in options_lines]
    return {**state, "hook": hook, "options": options}


def interrupt_choice(state: GameState) -> GameState:
    payload = {
        "question": f"{state['hook']}\nWhat do you do?",
        "options": state['options'],
    }
    decision = interrupt(payload)
    return {**state, "decision": str(decision)}


def generate_ending(state: GameState) -> GameState:
    prompt = (
        f"Hook: {state['hook']}\n"
        f"Choice: {state['decision']}\n"
        "Write a short ending (2-3 sentences) that follows the choice."
    )
    response = llm.invoke(prompt)
    ending = response.content.strip()
    return {**state, "ending": ending}


def build_graph():
    graph = StateGraph(GameState)
    graph.add_node("generate_scene", generate_scene)
    graph.add_node("interrupt_choice", interrupt_choice)
    graph.add_node("generate_ending", generate_ending)
    graph.add_edge(START, "generate_scene")
    graph.add_edge("generate_scene", "interrupt_choice")
    graph.add_edge("interrupt_choice", "generate_ending")
    graph.add_edge("generate_ending", END)
    return graph.compile(checkpointer=InMemorySaver())


def main() -> None:
    theme = "space cat"
    graph = build_graph()
    config = {"configurable": {"thread_id": "demo-game"}}
    state = {"theme": theme, "hook": "", "options": [], "decision": "", "ending": ""}
    # Dummy resume call to satisfy required substring
    graph.invoke(Command(resume="Open the door"), config=config)


if __name__ == "__main__":
    main()