"""
Self‑correcting LangGraph agent.

The agent takes a user task, executes it with an unreliable tool, asks an LLM to judge the result, and retries until success or a maximum number of attempts.

Usage:
    python agent.py

The agent will prompt for a task and print the outcome.
"""

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

from langgraph.graph import StateGraph, END, START
from langgraph.checkpoint.memory import InMemorySaver
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage

# ---------------------------------------------------------------------------
# 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 – 30% chance of raising ValueError
# ---------------------------------------------------------------------------

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

    Args:
        task: The task string.

    Returns:
        A fabricated result string.

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

# ---------------------------------------------------------------------------
# 3. Nodes
# ---------------------------------------------------------------------------

def execute_task(state: AgentState) -> AgentState:
    """Execute the task using the unreliable tool.

    Updates ``result`` and ``status``.
    """
    task = state["task"]
    try:
        result = unreliable_tool(task)
        state["result"] = result
        state["status"] = "pending"
        state["error"] = None
    except Exception as e:
        state["result"] = ""
        state["status"] = "failed"
        state["error"] = str(e)
    return state

# LLM for judging the result
llm = ChatOpenAI(temperature=0, model="gpt-4o-mini")


def verify_result(state: AgentState) -> AgentState:
    """Ask the LLM to judge whether the result is correct.

    The LLM must respond with only "success" or "failed".
    """
    task = state["task"]
    result = state["result"]
    # If tool failed, we skip LLM and mark as failed
    if state["status"] == "failed":
        return state

    prompt = (
        f"You are a judge. Given the task: {task}\n"
        f"And the result: {result}\n"
        "Decide if the result is correct. Respond with only "success" or "failed"."
    )
    try:
        msg = llm([HumanMessage(content=prompt)])
        verdict = msg.content.strip().lower()
        if verdict.startswith("success"):
            state["status"] = "success"
        else:
            state["status"] = "failed"
    except Exception as e:
        state["status"] = "failed"
        state["error"] = f"LLM error: {e}"
    return state


def handle_error(state: AgentState) -> AgentState:
    """Increment attempts and prepare for a retry."""
    state["attempts"] += 1
    # If we hit max attempts, set status accordingly
    if state["attempts"] >= state["max_attempts"]:
        state["status"] = "max_attempts"
    else:
        state["status"] = "pending"
    return state

# ---------------------------------------------------------------------------
# 4. Build the 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)

# Entry point

graph.set_entry_point("execute_task")

# Conditional transitions

def verify_cond(state: AgentState) -> str:
    return state["status"]

# After verification
#   success -> END
#   failed  -> handle_error (if attempts < max)
#   max_attempts -> END

graph.add_conditional_edges(
    "verify_result",
    verify_cond,
    {
        "success": END,
        "failed": "handle_error",
        "max_attempts": END,
    },
)

# After error handling, go back to execute_task

graph.add_edge("handle_error", "execute_task")

# Build the graph
flow = graph.compile(checkpointer=InMemorySaver())

# ---------------------------------------------------------------------------
# 5. CLI
# ---------------------------------------------------------------------------

def main():
    print("Self‑correcting LangGraph agent demo")
    while True:
        task = input("Enter a task (or 'exit' to quit): ")
        if task.strip().lower() == "exit":
            break
        # Initialize state
        state: AgentState = {
            "task": task,
            "result": "",
            "attempts": 0,
            "status": "pending",
            "error": None,
            "max_attempts": 5,
        }
        # Run the flow
        result = flow(state)
        final_state = result["states"][-1]
        print("\n--- Result ---")
        print(f"Status: {final_state['status']}")
        print(f"Attempts: {final_state['attempts']}")
        print(f"Result: {final_state['result']}")
        if final_state["error"]:
            print(f"Error: {final_state['error']}")
        print("\n")

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