"""
Self‑correcting LangGraph agent demo.

Run:
    python main.py

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

import random
from typing import TypedDict, Optional

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

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

# ---------------------------------------------------------------------------
# 2. Unreliable tool
# ---------------------------------------------------------------------------

def unreliable_addition(a: int, b: int) -> str:
    """Adds two numbers but fails with 30 % probability.

    Returns the sum as a string.
    """
    if random.random() < 0.3:
        raise ValueError("Random failure in unreliable tool")
    return str(a + b)

# ---------------------------------------------------------------------------
# 3. LLM initialization
# ---------------------------------------------------------------------------
# The OpenAI key must be set in the environment.
llm = ChatOpenAI(temperature=0)

# ---------------------------------------------------------------------------
# 4. Node implementations
# ---------------------------------------------------------------------------

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

    The task is expected to be a string of the form "<a> + <b>".
    """
    # Increment attempt counter at the start of each execution
    state["attempts"] += 1

    task = state["task"].strip()
    try:
        a_str, b_str = task.split("+")
        a = int(a_str.strip())
        b = int(b_str.strip())
        result = unreliable_addition(a, b)
        state["result"] = result
        state["error"] = None
    except Exception as e:
        state["result"] = ""
        state["error"] = str(e)
    return state

def verify_result(state: AgentState) -> AgentState:
    """LLM checks if the result is correct.

    The LLM is asked to output only "success" or "failed".
    """
    task = state["task"].strip()
    result = state["result"].strip()
    # Compute expected result for comparison
    try:
        a_str, b_str = task.split("+")
        expected = str(int(a_str.strip()) + int(b_str.strip()))
    except Exception:
        expected = ""

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

def handle_error(state: AgentState) -> AgentState:
    """Decide whether to retry or stop based on the number of attempts."""
    if state["attempts"] >= state["max_attempts"]:
        state["status"] = "max_attempts"
    else:
        state["status"] = "pending"
    return state

# ---------------------------------------------------------------------------
# 5. Graph construction
# ---------------------------------------------------------------------------

graph = StateGraph(AgentState)

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.add_edge(START, "execute_task")
# execute_task -> verify_result
graph.add_edge("execute_task", "verify_result")
# verify_result -> END on success, else -> handle_error on failed

graph.add_conditional_edges(
    "verify_result",
    lambda state: "END" if state["status"] == "success" else "handle_error" if state["status"] == "failed" else "END",
)
# handle_error -> execute_task if pending, else END

graph.add_conditional_edges(
    "handle_error",
    lambda state: "execute_task" if state["status"] == "pending" else "END",
)

# Compile the graph
app = graph.compile()

# ---------------------------------------------------------------------------
# 6. Demo execution
# ---------------------------------------------------------------------------

def main():
    task = "2 + 2"
    initial_state: AgentState = {
        "task": task,
        "result": "",
        "attempts": 0,
        "status": "pending",
        "error": None,
        "max_attempts": 5,
    }

    print(f"Task: {task}")

    # Run the graph
    for event in app.stream(initial_state):
        state = event["state"]
        attempts = state["attempts"]
        if state["status"] == "pending":
            print(f"Попытка {attempts}: выполняется...")
        elif state["status"] == "failed":
            print(f"Попытка {attempts}: ошибка -> {state['error']}")
        elif state["status"] == "success":
            print(f"Попытка {attempts}: результат {state['result']} → verify: success")
        elif state["status"] == "max_attempts":
            print(f"Попытка {attempts}: достигнут лимит попыток → verify: failed")

    final_state = app.get_state()
    print("\nИтог:")
    if final_state["status"] == "success":
        print(f"success за {final_state['attempts']} попыток")
    else:
        print(f"failed после {final_state['attempts']} попыток")

if __name__ == "__main__":
    main()