from __future__ import annotations

from typing import TypedDict

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


class GameState(TypedDict):
    scene: str
    inventory: list[str]
    decision: str


def intro(state: GameState) -> GameState:
    return {
        "scene": "Вы стоите у двери лаборатории AI-агента.",
        "inventory": state.get("inventory", []),
        "decision": "",
    }


def ask_human(state: GameState) -> GameState:
    decision = interrupt(
        {
            "scene": state["scene"],
            "question": "Открыть дверь или осмотреть коридор?",
            "options": ["open", "look"],
        }
    )
    return {**state, "decision": str(decision)}


def resolve(state: GameState) -> GameState:
    inventory = list(state.get("inventory", []))
    if state.get("decision") == "open":
        inventory.append("access_card")
        scene = "Дверь открыта, агент получил access_card."
    else:
        scene = "Вы нашли подсказку на стене и вернулись к двери."
    return {**state, "scene": scene, "inventory": inventory}


def build_graph():
    graph = StateGraph(GameState)
    graph.add_node("intro", intro)
    graph.add_node("ask_human", ask_human)
    graph.add_node("resolve", resolve)
    graph.add_edge(START, "intro")
    graph.add_edge("intro", "ask_human")
    graph.add_edge("ask_human", "resolve")
    graph.add_edge("resolve", END)
    return graph.compile(checkpointer=InMemorySaver())


def main() -> None:
    graph = build_graph()
    config = {"configurable": {"thread_id": "demo-game"}}
    print(graph.invoke({"scene": "", "inventory": [], "decision": ""}, config=config))
    print("Resume example:")
    print(graph.invoke(Command(resume="open"), config=config))  # type: ignore[name-defined]


if __name__ == "__main__":
    main()