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 ----------
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 ----------
backend = CompositeBackend(
    [
        LocalShellBackend(workspace_dir="./workspace"),
        FilesystemBackend(),
    ]
)

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


# ---------- DeepAgent (used inside execute_task node) ----------
deep_agent = create_deep_agent(
    model=llm,
    tools=[unreliable_tool],
    backend=backend,
    system_prompt="You are a helpful assistant that uses the provided tool to answer user queries.",
)

# ---------- State definition ----------
class AgentState(TypedDict):
    task: str
    result: str
    attempts: int
    status: str  # pending | success | failed | max_attempts
    error: str | None
    max_attempts: int
    messages: Annotated[list, add_messages]

# ---------- Nodes ----------
async def execute_task(state: AgentState):
    """Run the task using the deep agent."""
    try:
        response = await deep_agent.ainvoke(
            {"messages": [HumanMessage(content=state["task"])]},
            {"configurable": {"thread_id": f"session-{state['attempts']}"}},
        )
        # The deep agent returns a dict with "messages"
        result_msg = response["messages"][-1].content
        return {
            "result": result_msg,
            "error": None,
            "status": "pending",
            "messages": response["messages"],
        }
    except Exception as e:
        return {
            "result": "",
            "error": str(e),
            "status": "failed",
            "messages": [],
        }


async def verify_result(state: AgentState):
    """Ask LLM to judge the result."""
    judge_prompt = f"""You are a judge. Determine if the following result correctly solves the task.

Task: {state['task']}
Result: {state['result']}

Respond with only one word: SUCCESS if the result is correct, otherwise FAILED."""
    judge_response = await llm.ainvoke([HumanMessage(content=judge_prompt)])
    verdict = judge_response.content.strip().lower()
    if verdict == "success":
        new_status = "success"
    else:
        new_status = "failed"
    return {"status": new_status, "messages": [HumanMessage(content=judge_response.content)]}


def handle_error(state: AgentState):
    """Increase attempt counter and decide next step."""
    attempts = state["attempts"] + 1
    if attempts >= state["max_attempts"]:
        return {
            "attempts": attempts,
            "status": "max_attempts",
            "error": state.get("error"),
        }
    else:
        return {
            "attempts": attempts,
            "status": "pending",
            "error": None,
        }

# ---------- Graph ----------
graph = StateGraph(AgentState)

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

graph.add_edge(START, "execute_task")
graph.add_edge("execute_task", "verify_result")
graph.add_conditional_edges(
    "verify_result",
    lambda state: state["status"],
    {
        "success": END,
        "failed": "handle_error",
        "max_attempts": END,
    },
)
graph.add_edge("handle_error", "execute_task")

graph.set_entry_point(START)

app = graph.compile()

# ---------- Main ----------
async def main():
    task_description = "Calculate 2+2."
    initial_state: AgentState = {
        "task": task_description,
        "result": "",
        "attempts": 0,
        "status": "pending",
        "error": None,
        "max_attempts": 5,
        "messages": [],
    }

    async for event in app.astream(initial_state):
        # Print progress information
        if "attempts" in event:
            print(f"Attempt {event['attempts']}: status={event['status']}")
        if event.get("error"):
            print(f"Error: {event['error']}")
        if event.get("result"):
            print(f"Result: {event['result']}")

    final = await app.ainvoke(initial_state)
    print("\n=== Final Outcome ===")
    print(f"Task: {task_description}")
    print(f"Status: {final['status']}")
    print(f"Attempts: {final['attempts']}")
    if final["status"] == "success":
        print(f"Successful result: {final['result']}")
    else:
        print("Failed to obtain a correct result.")

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