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

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

# ---------- Example tool ----------
@tool
def echo_tool(text: str) -> str:
    """Return the same text back."""
    return text

# ---------- DeepAgent ----------
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 interrupt ----------
def ask_human_node(state: GraphState):
    # Prepare payload for interrupt
    payload = {
        "type": "confirm",
        "question": "Do you want to continue?",
        "options": ["approve", "reject"],
    }
    # Raise interrupt; execution will pause here
    return interrupt(payload)

# ---------- Node after resume ----------
def after_human_node(state: GraphState):
    # The resumed payload will contain the answer under key 'answer'
    answer = state.get("human_value", "no answer")
    # Use deepagent to produce a final message
    result = asyncio.run(
        agent.ainvoke(
            {"messages": [HumanMessage(content=f"User answered: {answer}")]},
            {"configurable": {"thread_id": "deepagent-session"}},
        )
    )
    # Append the agent's response to the message list
    state["messages"].append(result["messages"][-1])
    return state

# ---------- Build graph ----------
graph = StateGraph(GraphState)

graph.add_node("ask_human", ask_human_node)
graph.add_node("after_human", after_human_node)

graph.add_edge(START, "ask_human")
graph.add_edge("ask_human", "after_human")
graph.add_edge("after_human", END)

# Use in-memory checkpointing
graph.set_checkpoint_saver(InMemorySaver())
app = graph.compile()

# ---------- Runtime loop handling interrupt ----------
async def run():
    thread_id = "example-thread"
    config = {"configurable": {"thread_id": thread_id}}

    # Initial stream
    stream = app.stream(
        {"messages": []},
        config,
    )

    async for chunk in stream:
        # Check for interrupt signal
        if "__interrupt__" in chunk:
            interrupt_payload = chunk["__interrupt__"][0].value
            print("\n--- Interrupt received ---")
            print(f"Type: {interrupt_payload.get('type')}")
            print(f"Question: {interrupt_payload.get('question')}")
            # Simple console input (could use questionary)
            while True:
                answer = input(f"Choose {interrupt_payload.get('options')}: ").strip()
                if answer in interrupt_payload.get("options"):
                    break
                print("Invalid option, try again.")
            # Add answer to payload
            interrupt_payload["answer"] = answer

            # Resume graph with the updated payload
            resume_cmd = Command(resume=interrupt_payload)
            resume_stream = app.stream(resume_cmd, config)

            async for resume_chunk in resume_stream:
                if "__interrupt__" in resume_chunk:
                    # Should not happen in this simple example
                    continue
                if "messages" in resume_chunk:
                    # Final state reached
                    final_state = resume_chunk
                    print("\n--- Final state ---")
                    for msg in final_state["messages"]:
                        print(msg.content)
                    return

# ---------- Entry point ----------
if __name__ == "__main__":
    asyncio.run(run())