from typing import TypedDict, Optional

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


class GraphState(TypedDict):
    human_value: Optional[str] = None


llm = ChatOpenAI(
    model="openai/gpt-oss-20b",
    base_url='https://platform.brojs.ru/jrnl-bh/api/inference/v1',
    api_key=SecretStr("jrnl_30283ab953615cbb6846ff9940a1eedce0b76d7b2f59a2394f29e74643e6a90d"),
    temperature=0.7,
)


def ask(state: GraphState) -> Command:
    # Trigger an interrupt with a question and options
    return interrupt(
        type="question",
        question="Choose your favorite color:",
        options=["red", "green", "blue"],
    )


# Build the graph
builder = StateGraph(GraphState)
builder.add_node("ask", ask)

# Set entry point correctly (only one argument – node name)
builder.set_entry_point("ask")

# Define edges
builder.add_edge(START, "ask")
builder.add_edge("ask", END)

# Add a checkpoint
checkpoint = InMemorySaver()
graph = builder.compile(checkpointer=checkpoint)


def main():
    state: GraphState = {"human_value": None}
    while True:
        result = graph.invoke(state)
        # If the graph finished without interruption
        if "messages" in result and not any(
            getattr(msg, "type", "") == "__interrupt__" for msg in result["messages"]
        ):
            print("Graph finished. Result:", state)
            break

        # Handle interrupt
        for msg in result["messages"]:
            if getattr(msg, "type", "") == "__interrupt__":
                question = msg.content.get("question")
                options = msg.content.get("options")
                print(question)
                for i, opt in enumerate(options, 1):
                    print(f"{i}. {opt}")
                choice_idx = int(input("Select option number: ")) - 1
                state["human_value"] = options[choice_idx]
                # Resume the graph with updated state
                # The resume is implicit by invoking again after setting the value
                break


if __name__ == "__main__":
    main()