#!/usr/bin/env python3
"""Self‑correcting LangGraph agent.

Run with:
    python main.py "Вычисли 2+2"
"""
import os
import random
import sys
from typing import TypedDict, Dict, Any

from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import InMemorySaver

# ---------- LLM ----------
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://openrouter.ai/api/v1",
    api_key=os.getenv("OPENAI_API_KEY"),
    temperature=0.0,
)

# ---------- State ----------
class AgentState(TypedDict):
    task: str
    result: str
    attempts: int
    status: str  # pending | success | failed | max_attempts
    error: str | None
    max_attempts: int

# ---------- Tool ----------
class UnreliableTool:
    """Tool that fails ~30% of the time."""

    def __call__(self, input: str) -> str:
        if random.random() < 0.3:
            raise ValueError("Simulated tool failure")
        # Simple arithmetic evaluator for demo purposes
        try:
            return str(eval(input))
        except Exception as e:
            raise ValueError(f"Evaluation error: {e}")

unreliable_tool = UnreliableTool()

# ---------- Nodes ----------
async def execute_task(state: AgentState) -> Dict[str, Any]:
    """Run the task using the unreliable tool."""
    try:
        result = unreliable_tool(state["task"])
        return {"result": result, "error": None, "status": "pending"}
    except Exception as e:
        return {"result": "", "error": str(e), "status": "failed"}

async def verify_result(state: AgentState) -> Dict[str, Any]:
    """LLM judge: success or failed."""
    prompt = (
        f"Task: {state['task']}\n"
        f"Result: {state['result']}\n"
        f"Error: {state['error']}\n"
        "Is the result correct? Answer with only 'success' or 'failed'."
    )
    response = await llm.ainvoke(prompt)
    verdict = response.content.strip().lower()
    if verdict not in {"success", "failed"}:
        verdict = "failed"
    return {"status": verdict}

async def handle_error(state: AgentState) -> Dict[str, Any]:
    """Increment attempts and decide whether to retry."""
    attempts = state["attempts"] + 1
    if attempts >= state["max_attempts"]:
        return {"attempts": attempts, "status": "max_attempts"}
    return {"attempts": attempts, "status": "pending", "error": None, "result": ""}

# ---------- 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")
builder.add_conditional_edges(
    "verify_result",
    lambda x: x["status"],
    {
        "success": END,
        "failed": "handle_error",
        "max_attempts": END,
    },
)
builder.add_edge("handle_error", "execute_task")

graph = builder.compile(checkpointer=InMemorySaver())

# ---------- Runner ----------
async def run(task: str, max_attempts: int = 5):
    initial_state: AgentState = {
        "task": task,
        "result": "",
        "attempts": 0,
        "status": "pending",
        "error": None,
        "max_attempts": max_attempts,
    }
    async for event in graph.astream_events(initial_state, version="1"):
        if "update" in event:
            state = event["update"]
            if state["status"] == "pending" and state["attempts"] > 0:
                print(f"Попытка {state['attempts']}: {state['error'] or state['result']}")
    final_state = graph.get_state()
    print("\nИтог:", final_state["status"], f"за {final_state['attempts']} попытки(и)")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python main.py '<task>'")
        sys.exit(1)
    task = sys.argv[1]
    import asyncio
    asyncio.run(run(task))