"""Self‑correcting LangGraph agent.

The agent receives a *task* string. It executes the task via an
`unreliable_tool` that sometimes raises a ValueError. After execution
it asks the LLM (OpenAI or Ollama) to judge whether the *result* is
correct. If the judge says ``failed`` the agent retries until
``max_attempts`` is reached.

The implementation uses LangGraph's low‑level API: a StateGraph
with three nodes – execute_task, verify_result, handle_error – and a
simple loop.

Run the script with:

    python agent.py

It will ask for a task, then show the attempts and final status.
"""

from __future__ import annotations

import random
import sys
from typing import TypedDict

from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END

# ---------------------------------------------------------------------------
# 1. State definition
# ---------------------------------------------------------------------------

class AgentState(TypedDict):
    task: str
    result: str | None
    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 sometimes fails.

    The function simply returns ``task`` reversed (as a dummy result) but
    raises a ValueError 30 % of the time.
    """
    if random.random() < 0.3:
        raise ValueError("Simulated tool failure")
    return task[::-1]  # dummy "computation"

# ---------------------------------------------------------------------------
# 3. LLM judge – asks for "success" or "failed"
# ---------------------------------------------------------------------------

llm = ChatOpenAI(temperature=0, model="gpt-4o-mini")  # or use Ollama

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

    The prompt forces the model to answer only "success" or "failed".
    """
    if state["result"] is None:
        # Should not happen – guard
        state["status"] = "failed"
        return state

    prompt = (
        f"Task: {state['task']}\n"
        f"Result: {state['result']}\n"
        "Is this result correct? Respond with only 'success' or 'failed'."
    )
    response = llm.invoke(prompt)
    verdict = response.content.strip().lower()
    if verdict == "success":
        state["status"] = "success"
    else:
        state["status"] = "failed"
    return state

# ---------------------------------------------------------------------------
# 4. Execute task node
# ---------------------------------------------------------------------------

async def execute_task(state: AgentState) -> AgentState:
    """Run the unreliable tool and capture errors."""
    try:
        result = unreliable_tool(state["task"])
        state["result"] = result
        state["error"] = None
    except Exception as exc:  # catch ValueError
        state["result"] = None
        state["error"] = str(exc)
    return state

# ---------------------------------------------------------------------------
# 5. Handle error / retry node
# ---------------------------------------------------------------------------

async def handle_error(state: AgentState) -> AgentState:
    """Increment attempt counter and decide whether to retry."""
    state["attempts"] += 1
    if state["attempts"] >= state["max_attempts"]:
        state["status"] = "max_attempts"
    else:
        # Reset result and error for next try
        state["result"] = None
        state["error"] = None
    return state

# ---------------------------------------------------------------------------
# 6. Build the graph
# ---------------------------------------------------------------------------

graph = StateGraph(AgentState)

# Add nodes
graph.add_node("execute_task", execute_task)
graph.add_node("verify_result", verify_result)
graph.add_node("handle_error", handle_error)

# Define edges
# Start → execute_task
graph.set_entry_point("execute_task")

# After execution, go to verification
graph.add_edge("execute_task", "verify_result")

# Verification outcomes
# success → END
# failed → check attempts
# max_attempts → END

# We use a conditional edge on the status field

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

# Map status to next node
graph.add_conditional_edges(
    "verify_result",
    verify_cond,
    {
        "success": END,
        "failed": "handle_error",
        "max_attempts": END,
    },
)

# From handle_error back to execute_task
graph.add_edge("handle_error", "execute_task")

# Compile the graph into a runnable chain
agent = graph.compile()

# ---------------------------------------------------------------------------
# 7. CLI entry point
# ---------------------------------------------------------------------------

def main() -> None:
    print("Self‑correcting LangGraph agent demo")
    task = input("Enter a task: ")
    if not task:
        print("No task provided. Exiting.")
        sys.exit(0)

    # Initial state
    state: AgentState = {
        "task": task,
        "result": None,
        "attempts": 1,
        "status": "pending",
        "error": None,
        "max_attempts": 5,
    }

    # Run the chain
    final_state = agent.invoke(state)

    print("\n--- Result ---")
    print(f"Task: {final_state['task']}")
    print(f"Attempts: {final_state['attempts']}")
    print(f"Status: {final_state['status']}")
    if final_state['result']:
        print(f"Result: {final_state['result']}")
    if final_state['error']:
        print(f"Last error: {final_state['error']}")


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