import os
import asyncio
import random
from typing import TypedDict, Annotated

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

from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages

# ----------------------------------------------------------------------
# LLM configuration (OpenRouter, free tier)
# ----------------------------------------------------------------------
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 (required by the framework)
# ----------------------------------------------------------------------
backend = CompositeBackend(
    [
        LocalShellBackend(workspace_dir="./workspace"),
        FilesystemBackend(),
    ]
)

# ----------------------------------------------------------------------
# Unreliable tool used to demonstrate retry logic
# ----------------------------------------------------------------------
@tool
def unreliable_tool(query: str) -> str:
    """
    Simulates an unreliable external service.
    With ~30% probability it raises a ValueError.
    """
    if random.random() < 0.3:
        raise ValueError("Simulated tool failure")
    return f"Result for '{query}'"

# ----------------------------------------------------------------------
# DeepAgent creation (required by the course)
# ----------------------------------------------------------------------
deep_agent = create_deep_agent(
    model=llm,
    tools=[unreliable_tool],
    backend=backend,
    system_prompt="You are a helpful assistant that can use tools when needed.",
)

# ----------------------------------------------------------------------
# State definition for the LangGraph workflow
# ----------------------------------------------------------------------
class AgentState(TypedDict):
    task: str
    result: str
    attempts: int
    status: str  # pending | success | failed | max_attempts
    error: str | None
    max_attempts: int

# ----------------------------------------------------------------------
# Node: execute_task
# Calls the deep agent to perform the task using the unreliable tool.
# ----------------------------------------------------------------------
async def execute_task(state: AgentState) -> AgentState:
    try:
        # Invoke the deep agent with the current task description
        response = await deep_agent.ainvoke(
            {"messages": [HumanMessage(content=state["task"])]},
            {"configurable": {"thread_id": f"session-{state['attempts'] + 1}"}},
        )
        # Extract the assistant's final message
        result_msg = response["messages"][-1].content
        state["result"] = result_msg
        state["error"] = None
    except Exception as e:
        # Capture any exception from the tool or agent
        state["result"] = ""
        state["error"] = str(e)
    return state

# ----------------------------------------------------------------------
# Node: verify_result
# Uses the LLM as a judge to decide if the result is acceptable.
# ----------------------------------------------------------------------
async def verify_result(state: AgentState) -> AgentState:
    # Prompt the LLM to judge the result. We ask for a strict "success" or "failed".
    judge_prompt = (
        "You are a verifier. Given the original task and the agent's result, "
        "respond with only the word 'success' if the result correctly fulfills the task, "
        "otherwise respond with 'failed'. Do not add any other text."
    )
    messages = [
        HumanMessage(content=judge_prompt),
        HumanMessage(content=f"Task: {state['task']}\nResult: {state['result']}"),
    ]
    judge_response = await llm.ainvoke(messages)
    verdict = judge_response.content.strip().lower()
    if verdict == "success":
        state["status"] = "success"
    else:
        state["status"] = "failed"
    return state

# ----------------------------------------------------------------------
# Node: handle_error
# Increments attempts and decides whether to retry or stop.
# ----------------------------------------------------------------------
def handle_error(state: AgentState) -> AgentState:
    state["attempts"] += 1
    if state["attempts"] >= state["max_attempts"]:
        state["status"] = "max_attempts"
    else:
        state["status"] = "pending"
    return state

# ----------------------------------------------------------------------
# Build the StateGraph with the defined nodes and transitions
# ----------------------------------------------------------------------
workflow = StateGraph(AgentState)

workflow.add_node("execute_task", execute_task)
workflow.add_node("verify_result", verify_result)
workflow.add_node("handle_error", handle_error)

workflow.add_edge(START, "execute_task")
workflow.add_edge("execute_task", "verify_result")
workflow.add_conditional_edges(
    "verify_result",
    lambda state: "success" if state["status"] == "success" else "retry",
    {
        "success": END,
        "retry": "handle_error",
    },
)
workflow.add_edge("handle_error", "execute_task")
workflow.add_conditional_edges(
    "handle_error",
    lambda state: "end" if state["status"] in ("max_attempts", "success") else "retry",
    {
        "end": END,
        "retry": "execute_task",
    },
)

graph = workflow.compile()

# ----------------------------------------------------------------------
# Main entry point: runs the graph for a sample task and prints progress
# ----------------------------------------------------------------------
async def main():
    initial_state: AgentState = {
        "task": "Calculate 2+2 and return the answer as a plain number.",
        "result": "",
        "attempts": 0,
        "status": "pending",
        "error": None,
        "max_attempts": 5,
    }

    async for event in graph.astream(initial_state):
        # The graph yields intermediate states; we log useful info.
        if "attempts" in event:
            print(f"Attempt {event['attempts']}: status={event['status']}", end="")
            if event["error"]:
                print(f", error={event['error']}")
            else:
                print(f", result={event['result'][:50]}")
        if event["status"] in ("success", "max_attempts"):
            print("\nFinal status:", event["status"])
            print("Result:", event["result"])
            break

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