"""
Self‑correcting LangGraph agent.

The agent:
1. Takes a user task.
2. Executes it via an unreliable tool.
3. Asks an LLM to judge the result (success / failed).
4. Retries until success or max_attempts.

Run with:
    python agent.py

Requires an OpenAI API key in the environment variable `OPENAI_API_KEY`.
"""

from __future__ import annotations

import os
import random
import time
from typing import Dict, TypedDict

from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode
from langgraph.checkpoint.memory import InMemorySaver

# ---------------------------------------------------------------------------
# 1. State definition
# ---------------------------------------------------------------------------
class AgentState(TypedDict):
    task: str
    result: str
    attempts: int
    status: str  # pending | success | failed | max_attempts
    error: str | None
    max_attempts: int

# ---------------------------------------------------------------------------
# 2. Unreliable tool
# ---------------------------------------------------------------------------
def unreliable_tool(task: str) -> str:
    """Simulates a tool that fails 30% of the time.

    Args:
        task: The task string.

    Returns:
        A string result (here we just echo the task for demo).

    Raises:
        ValueError: Simulated failure.
    """
    if random.random() < 0.3:
        raise ValueError("Simulated tool failure")
    # Simulate some work
    time.sleep(0.5)
    return f"Result for task: {task}"

# ---------------------------------------------------------------------------
# 3. LLM judge node
# ---------------------------------------------------------------------------
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

JUDGE_PROMPT = (
    "You are a judge that evaluates the result of a task. "
    "Given the task and the result, reply with either 'success' or 'failed'. "
    "Do not add any other text."
)

async def verify_result(state: AgentState) -> Dict[str, str]:
    """LLM judge that returns status.

    Returns a dict with key 'status' set to 'success' or 'failed'.
    """
    task = state["task"]
    result = state["result"]
    prompt = f"Task: {task}\nResult: {result}\n{JUDGE_PROMPT}"
    response = await llm.ainvoke(prompt)
    status = response.content.strip().lower()
    if status not in {"success", "failed"}:
        # Fallback to failed if LLM is uncertain
        status = "failed"
    return {"status": status}

# ---------------------------------------------------------------------------
# 4. Handle error / retry node
# ---------------------------------------------------------------------------
async def handle_error(state: AgentState) -> Dict[str, int | str | None]:
    """Increment attempts and decide whether to retry or stop.

    Returns updated attempts and status.
    """
    attempts = state["attempts"] + 1
    max_attempts = state["max_attempts"]
    status = "max_attempts" if attempts >= max_attempts else "pending"
    return {"attempts": attempts, "status": status}

# ---------------------------------------------------------------------------
# 5. Execute task node
# ---------------------------------------------------------------------------
async def execute_task(state: AgentState) -> Dict[str, str | None]:
    """Runs the unreliable tool and captures result or error.

    Returns updated result and error.
    """
    task = state["task"]
    try:
        result = unreliable_tool(task)
        return {"result": result, "error": None}
    except Exception as e:
        return {"result": "", "error": str(e)}

# ---------------------------------------------------------------------------
# 6. Build the graph
# ---------------------------------------------------------------------------
builder = StateGraph(AgentState)

# Register nodes
builder.add_node("execute_task", execute_task)
builder.add_node("verify_result", verify_result)
builder.add_node("handle_error", handle_error)

# Define edges
builder.set_entry_point("execute_task")

# After executing, verify
builder.add_edge("execute_task", "verify_result")

# After verification
builder.add_conditional_edges(
    "verify_result",
    lambda x: x["status"],
    {
        "success": END,
        "failed": "handle_error",
        "max_attempts": END,  # safety, though not expected here
    },
)

# After error handling, either retry or end
builder.add_conditional_edges(
    "handle_error",
    lambda x: x["status"],
    {
        "pending": "execute_task",
        "max_attempts": END,
    },
)

# Compile graph
graph = builder.compile(checkpointer=InMemorySaver())

# ---------------------------------------------------------------------------
# 7. CLI
# ---------------------------------------------------------------------------
def main() -> None:
    print("Self‑correcting LangGraph agent demo")
    task = input("Enter a task: ")
    max_attempts = 5

    initial_state: AgentState = {
        "task": task,
        "result": "",
        "attempts": 0,
        "status": "pending",
        "error": None,
        "max_attempts": max_attempts,
    }

    # Run the graph
    result = graph.invoke(initial_state)

    # Extract final status
    final_status = result["status"]
    attempts = result["attempts"]
    final_result = result["result"]
    error = result["error"]

    print("\n--- Result ---")
    print(f"Status: {final_status}")
    print(f"Attempts: {attempts}")
    if error:
        print(f"Last error: {error}")
    print(f"Result: {final_result}")


if __name__ == "__main__":
    main()
"