import os
import asyncio
from typing import TypedDict, Annotated

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

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 FilesystemBackend, LocalShellBackend, CompositeBackend

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

# ---------- State ----------
class GraphState(TypedDict):
    human_value: Annotated[str | None, "value chosen by user"]
    foo: Annotated[str, "initial data"]

# ---------- Node with interrupt ----------
async def interrupt_node(state: GraphState) -> GraphState:
    # Trigger interrupt with structured payload
    payload = {
        "type": "confirm",
        "question": "Уверены, что хотите продолжить?",
        "allow_responds": ["approve", "reject"],
    }
    # The node will pause here; after resume, the payload will be updated with answer
    await interrupt(payload)
    # After resume, payload will contain 'answer'
    answer = payload.get("answer")
    state["human_value"] = answer
    return state

# ---------- Graph ----------
builder = StateGraph(GraphState)
builder.add_node("interrupt_node", interrupt_node)
builder.set_entry_point("interrupt_node")
builder.add_edge("interrupt_node", END)

checkpoint = InMemorySaver()
graph = builder.compile(checkpointer=checkpoint)

# ---------- Tool to run the graph ----------
@tool
def run_graph() -> str:
    """Run the LangGraph with human‑in‑the‑loop interrupt."""
    thread_id = "thread-1"
    config = {"configurable": {"thread_id": thread_id}}
    # Initial state
    state: GraphState = {"human_value": None, "foo": "initial"}
    # Start streaming
    stream = graph.stream(state, config)
    try:
        for chunk in stream:
            if "__interrupt__" in chunk:
                # Extract payload
                payload = chunk["__interrupt__"][0].value
                # Show question to user
                answer = questionary.select(
                    payload["question"],
                    choices=payload["allow_responds"],
                ).ask()
                # Attach answer and resume
                payload["answer"] = answer
                stream = graph.stream(Command(resume=payload), config)
                continue
            # When stream ends, return final state
            if "node" in chunk:
                return f"Final state: {chunk['node']}"
    except Exception as e:
        return f"Error: {e}"
    return "Graph finished"

# ---------- DeepAgent ----------
backend = CompositeBackend([
    LocalShellBackend(workspace_dir="./workspace"),
    FilesystemBackend(),
])

agent = create_deep_agent(
    model=llm,
    tools=[run_graph],
    backend=backend,
    system_prompt="You are a helpful agent that can run a graph with human interrupt.",
)

async def main():
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content="run_graph")]},
        {"configurable": {"thread_id": "session-1"}},
    )
    print(result["messages"][-1].content)

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