"""
Self‑correcting LangGraph agent.

The agent receives a natural‑language task, executes it via an unreliable tool,
then asks an LLM to judge whether the result is correct. If the judge says
"failed" the agent retries until success or a maximum number of attempts.

The implementation uses LangGraph 1.x and LangChain 1.x.
"""

from __future__ import annotations

import random
import sys
from typing import TypedDict

# LangChain imports
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage

# LangGraph imports
from langgraph.graph import StateGraph, END

# ---------------------------------------------------------------------------
# 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:
    """Simulate a tool that fails 30 % of the time.

    The tool simply returns the string ``f"Result of {task}"`` but raises a
    ``ValueError`` with 30 % probability.
    """
    if random.random() < 0.3:
        raise ValueError("Simulated tool failure")
    return f"Result of {task}"

# ---------------------------------------------------------------------------
# 3. LLM judge
# ---------------------------------------------------------------------------

# Create a lightweight LLM instance.  The user must set the OPENAI_API_KEY
# environment variable or provide a key in the code.
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

# The judge prompt asks the model to answer only "success" or "failed".
JUDGE_PROMPT = (
    "You are a strict judge. Given the following result of a task, answer only "
    "one word: success or failed.\n\nResult: {result}\nAnswer:"  # no extra formatting
)

# ---------------------------------------------------------------------------
# 4. Node functions
# ---------------------------------------------------------------------------

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

    On success, store the result.  On failure, capture the exception.
    """
    try:
        result = unreliable_tool(state["task"])
        state["result"] = result
        state["error"] = None
    except Exception as exc:
        state["result"] = ""
        state["error"] = str(exc)
    return state

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

    The LLM must return either "success" or "failed".
    """
    if state["error"]:
        # If the tool raised an exception, we consider it a failure.
        state["status"] = "failed"
        return state

    # Build the prompt with the result.
    prompt = JUDGE_PROMPT.format(result=state["result"])
    messages = [HumanMessage(content=prompt)]
    ai_msg: AIMessage = await llm.ainvoke(messages)
    verdict = ai_msg.content.strip().lower()
    if verdict == "success":
        state["status"] = "success"
    else:
        state["status"] = "failed"
    return state

async def handle_error(state: AgentState) -> AgentState:
    """Increment attempts and prepare for retry.

    If the maximum number of attempts is reached, set status to
    "max_attempts".
    """
    state["attempts"] += 1
    if state["attempts"] >= state["max_attempts"]:
        state["status"] = "max_attempts"
    else:
        # Reset result and error for the next attempt.
        state["result"] = ""
        state["error"] = None
    return state

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

def build_graph(max_attempts: int = 5) -> StateGraph[AgentState]:
    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 the flow: execute → verify → (success → END | failed → handle_error → execute)
    graph.add_edge("execute_task", "verify_result")
    graph.add_edge("handle_error", "execute_task")

    # Conditional router based on status after verification.
    def router(state: AgentState) -> str:
        return state["status"]

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

    graph.set_entry_point("execute_task")
    return graph

# ---------------------------------------------------------------------------
# 6. CLI driver
# ---------------------------------------------------------------------------

async def main():
    if len(sys.argv) > 1:
        task = " ".join(sys.argv[1:])
    else:
        task = input("Введите задачу: ")

    max_attempts = 5
    graph = build_graph(max_attempts)
    app = graph.compile()

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

    # Run the graph until it ends.
    async for partial_state in app.stream(state):
        # Print progress when attempts change.
        if partial_state["attempts"] != state["attempts"]:
            print(f"Попытка {partial_state['attempts']}:", end=" ")
            if partial_state["error"]:
                print(f"Error → {partial_state['error']}")
            else:
                print(f"результат {partial_state['result']}")
        state = partial_state

    # Final status
    print("\nИтог:")
    if state["status"] == "success":
        print(f"Успех за {state['attempts']} попыток. Результат: {state['result']}")
    elif state["status"] == "max_attempts":
        print(f"Не удалось за {state['attempts']} попыток. Последняя ошибка: {state['error']}")
    else:
        print(f"Не удалось. Последняя ошибка: {state['error']}")


if __name__ == "__main__":
    import asyncio

    asyncio.run(main())