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

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

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
from langchain.tools import tool

import questionary

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

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

# ---------- Example tool (not used in this task but required for agent creation) ----------
@tool
def echo_tool(text: str) -> str:
    """Return the same text back."""
    return text

# ---------- DeepAgent (required by the course) ----------
agent = create_deep_agent(
    model=llm,
    tools=[echo_tool],
    backend=backend,
    system_prompt="You are a helpful assistant.",
)

# ---------- Graph state ----------
class GraphState(TypedDict):
    messages: Annotated[List[Any], add_messages]
    human_value: str

# ---------- Node that triggers a custom interrupt ----------
def ask_node(state: GraphState):
    # If we are resumed, the payload will contain the answer
    if isinstance(state, dict) and "answer" in state:
        # Store the answer and finish
        return {
            "messages": state.get("messages", []),
            "human_value": state["answer"],
        }
    # Otherwise raise an interrupt with the question payload
    payload = {
        "type": "confirm",
        "question": "Do you want to continue the workflow?",
        "options": ["approve", "reject"],
    }
    return interrupt(payload)

# ---------- Build the graph ----------
graph_builder = StateGraph(GraphState)
graph_builder.add_node("ask", ask_node)
graph_builder.add_edge(START, "ask")
graph_builder.add_edge("ask", END)

graph = graph_builder.compile(checkpointer=InMemorySaver())

# ---------- Execution loop handling interrupts ----------
async def run_graph():
    thread_id = "demo-thread"
    config = {"configurable": {"thread_id": thread_id}}

    # Initial state
    state: GraphState = {"messages": [], "human_value": ""}

    # Helper to process a stream until it finishes or hits an interrupt
    async def process_stream(initial):
        async for chunk in graph.astream(initial, config):
            # Detect interrupt
            if "__interrupt__" in chunk:
                return chunk  # return the interrupt chunk
            # Detect final state (contains human_value)
            if "human_value" in chunk:
                print("Workflow finished. Final state:")
                print(chunk)
                return None
        return None

    # First run - will hit the interrupt
    interrupt_chunk = await process_stream(state)
    while interrupt_chunk:
        payload = interrupt_chunk["__interrupt__"][0].value
        print("\n--- Human in the loop ---")
        answer = questionary.select(
            payload["question"], choices=payload["options"]
        ).ask()
        # Add answer to payload for resumption
        payload["answer"] = answer
        # Resume the graph with the updated payload
        interrupt_chunk = await process_stream(Command(resume=payload))

# ---------- Main entry ----------
if __name__ == "__main__":
    asyncio.run(run_graph())