"""
Self‑correcting LangGraph agent demo.

Run with:
    python main.py

Requires:
    pip install langgraph langchain-openai
"""
import random
from typing import TypedDict, Dict

# ---------- 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

# ---------- Unreliable tool ----------
class UnreliableTool:
    def run(self, input_: str) -> str:
        if random.random() < 0.3:
            raise ValueError("Simulated tool failure")
        # simple eval for demo purposes
        try:
            return str(eval(input_))
        except Exception as e:
            raise ValueError(f"Eval error: {e}")

# ---------- LangGraph imports ----------
from langgraph.graph import StateGraph, END

from langchain_openai import ChatOpenAI

# ---------- Nodes ----------
def execute_task(state: AgentState) -> AgentState:
    tool = UnreliableTool()
    try:
        result = tool.run(state["task"])
        state["result"] = result
        state["error"] = None
    except Exception as e:
        state["result"] = None
        state["error"] = str(e)
    return state

def verify_result(state: AgentState) -> AgentState:
    # Ask LLM to judge success or failed based on result and error
    llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
    prompt = (
        f"Task: {state['task']}\n"
        f"Result: {state.get('result')}\n"
        f"Error: {state.get('error')}\n"
        "Respond with only 'success' or 'failed'."
    )
    resp = llm.invoke(prompt)
    verdict = resp.content.strip().lower()
    if verdict not in ("success", "failed"):
        verdict = "failed"
    state["verdict"] = verdict
    return state

def handle_error(state: AgentState) -> AgentState:
    state["attempts"] += 1
    if state["attempts"] >= state["max_attempts"]:
        state["status"] = "max_attempts"
    else:
        state["status"] = "pending"
    return state

# ---------- Graph ----------
builder = StateGraph(AgentState)
builder.add_node("execute_task", execute_task)
builder.add_node("verify_result", verify_result)
builder.add_node("handle_error", handle_error)

builder.set_entry_point("execute_task")
builder.add_edge("execute_task", "verify_result")
# From verify_result, branch on verdict
builder.add_conditional_edges(
    "verify_result",
    lambda state: state.get("verdict"),
    {
        "success": END,
        "failed": "handle_error",
    },
)
builder.add_edge("handle_error", "execute_task")

graph = builder.compile()

# ---------- Demo runner ----------
if __name__ == "__main__":
    initial_state: AgentState = {
        "task": "2+2",
        "result": None,
        "attempts": 0,
        "status": "pending",
        "error": None,
        "max_attempts": 5,
    }
    for attempt in range(1, initial_state["max_attempts"] + 1):
        print(f"Attempt {attempt}:")
        result = graph.invoke(initial_state)
        if result.get("verdict") == "success":
            print(f"Success: {result['result']}")
            break
        else:
            print(f"Failed, error: {result.get('error')}")
    else:
        print("Reached max attempts without success.")
""