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

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, SystemMessage
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

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

# ----------------------------------------------------------------------
# Unreliable tool used for demonstration
# ----------------------------------------------------------------------
@tool
def unreliable_tool(query: str) -> str:
    """
    Simulates an unreliable external tool.
    With ~30% probability it raises a ValueError to trigger a retry.
    """
    if random.random() < 0.3:
        raise ValueError("Simulated tool failure")
    # Simple evaluation for arithmetic expressions
    try:
        result = eval(query, {"__builtins__": {}})
    except Exception:
        result = f"cannot evaluate: {query}"
    return str(result)

# ----------------------------------------------------------------------
# DeepAgent - 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.",
)

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

# ----------------------------------------------------------------------
# Node: execute_task
# ----------------------------------------------------------------------
async def execute_task(state: AgentState) -> dict:
    """Run the task using the unreliable tool."""
    try:
        # Call the tool directly; deep_agent is not needed here
        tool_result = unreliable_tool(state["task"])
        new_state = {
            "result": tool_result,
            "error": None,
            "status": "pending",
        }
    except Exception as e:
        new_state = {
            "result": "",
            "error": str(e),
            "status": "failed",
        }
    return new_state

# ----------------------------------------------------------------------
# Node: verify_result
# ----------------------------------------------------------------------
async def verify_result(state: AgentState) -> dict:
    """Ask LLM to judge whether the result satisfies the task."""
    judge_prompt = f"""You are a judge. The original task is:
{state['task']}

The agent produced the following result:
{state['result']}

Respond with only one word: "success" if the result correctly solves the task,
otherwise respond with "failed"."""
    messages = [
        SystemMessage(content="You are an objective judge."),
        HumanMessage(content=judge_prompt),
    ]
    response = await llm.ainvoke(messages)
    verdict = response.content.strip().lower()
    if verdict == "success":
        new_status = "success"
    else:
        new_status = "failed"
    return {"status": new_status, "error": None if new_status == "success" else "Verification failed"}

# ----------------------------------------------------------------------
# Node: handle_error
# ----------------------------------------------------------------------
async def handle_error(state: AgentState) -> dict:
    """Increase attempt counter and decide whether to retry."""
    attempts = state["attempts"] + 1
    if attempts >= state["max_attempts"]:
        return {"attempts": attempts, "status": "max_attempts", "error": "Maximum attempts reached"}
    else:
        return {"attempts": attempts, "status": "pending", "error": None, "result": ""}

# ----------------------------------------------------------------------
# Build the StateGraph
# ----------------------------------------------------------------------
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: state["status"],
    {
        "success": END,
        "failed": "handle_error",
        "max_attempts": END,
    },
)
workflow.add_edge("handle_error", "execute_task")

graph = workflow.compile()

# ----------------------------------------------------------------------
# Main entry point
# ----------------------------------------------------------------------
async def run_task(task: str, max_attempts: int = 5):
    initial_state: AgentState = {
        "task": task,
        "result": "",
        "attempts": 0,
        "status": "pending",
        "error": None,
        "max_attempts": max_attempts,
        "messages": [],
    }
    async for event in graph.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 = event
    print("\n=== Final Outcome ===")
    print(f"Task: {task}")
    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__":
    # Example task: simple arithmetic
    example_task = "2 + 2"
    asyncio.run(run_task(example_task, max_attempts=5))