"""
Self‑correcting LangGraph agent using deepagents.

Requirements:
- Python 3.10+
- deepagents, langchain-openai, langgraph
- OPENAI_API_KEY env var pointing to an OpenRouter key

Run:
    python main.py
"""

import os
import random
import asyncio
from typing import TypedDict

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langchain.tools import tool

from deepagents import create_deep_agent
from deepagents.backends import CompositeBackend, LocalShellBackend, FilesystemBackend

from langgraph.graph import StateGraph, START, END

# ---------------------------------------------------------------------------
# 1. LLM configuration (OpenRouter)
# ---------------------------------------------------------------------------
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,
)

# ---------------------------------------------------------------------------
# 2. Unreliable tool – 30 % chance of raising ValueError
# ---------------------------------------------------------------------------
@tool
def unreliable_tool(query: str) -> str:
    """Simulates an unreliable external tool.
    30 % of the time it raises ValueError to trigger a retry.
    """
    if random.random() < 0.3:
        raise ValueError("Simulated tool failure")
    return f"{query}"

# ---------------------------------------------------------------------------
# 3. Deepagents backend and agent
# ---------------------------------------------------------------------------
backend = CompositeBackend([
    LocalShellBackend(workspace_dir="./workspace"),
    FilesystemBackend(),
])

agent = create_deep_agent(
    model=llm,
    tools=[unreliable_tool],
    backend=backend,
    system_prompt="You are a helpful agent. Use the provided tool to compute the answer.",
)

# ---------------------------------------------------------------------------
# 4. Graph state definition
# ---------------------------------------------------------------------------
class AgentState(TypedDict):
    task: str
    result: str
    attempts: int
    status: str  # pending | success | failed | max_attempts
    error: str | None
    max_attempts: int

# ---------------------------------------------------------------------------
# 5. Graph nodes
# ---------------------------------------------------------------------------
async def execute_task(state: AgentState) -> AgentState:
    """Execute the task via the deepagents agent.
    """
    attempt_num = state["attempts"] + 1
    print(f"Попытка {attempt_num}:")
    try:
        # Invoke the agent – it will call the unreliable_tool internally
        result = await agent.ainvoke(
            {"messages": [HumanMessage(content=state["task"])]},
            {"configurable": {"thread_id": "session-1"}},
        )
        # The agent returns a dict with a "messages" list
        output = result["messages"][-1].content
        state["result"] = output
        state["error"] = None
        print(f"  Result: {output}")
    except Exception as e:
        state["result"] = ""
        state["error"] = str(e)
        print(f"  Error: {state['error']}")
    return state

async def verify_result(state: AgentState) -> AgentState:
    """Ask the LLM to judge whether the result is correct.
    The LLM must answer only "success" or "failed".
    """
    prompt = (
        f"Please evaluate the following result for the task '{state['task']}'.\n"
        f"Respond with only 'success' or 'failed'.\n"
        f"Result: {state['result']}"
    )
    verification = await llm.ainvoke([HumanMessage(content=prompt)])
    verdict = verification["content"].strip().lower()
    print(f"  Verify: {verdict}")
    if verdict.startswith("success"):
        state["status"] = "success"
    else:
        state["status"] = "failed"
    return state

async def handle_error(state: AgentState) -> AgentState:
    """Increment attempts and decide whether to retry or stop.
    """
    state["attempts"] += 1
    if state["attempts"] >= state["max_attempts"]:
        state["status"] = "max_attempts"
    else:
        state["status"] = "pending"
    return state

# ---------------------------------------------------------------------------
# 6. Build the 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.add_edge(START, "execute_task")
builder.add_edge("execute_task", "verify_result")

# Conditional transition after verification

def check_status(state: AgentState):
    if state["status"] == "success":
        return "success"
    if state["attempts"] >= state["max_attempts"]:
        return "max_attempts"
    return "handle_error"

builder.add_conditional_edges(
    "verify_result",
    check_status,
    {
        "success": "success",
        "max_attempts": "max_attempts",
        "handle_error": "handle_error",
    },
)

builder.add_edge("handle_error", "execute_task")
builder.add_edge("max_attempts", END)
builder.add_edge("success", END)

graph = builder.compile()

# ---------------------------------------------------------------------------
# 7. Demo run
# ---------------------------------------------------------------------------
async def main():
    task = "Compute 2+2"
    initial_state: AgentState = {
        "task": task,
        "result": "",
        "attempts": 0,
        "status": "pending",
        "error": None,
        "max_attempts": 5,
    }
    final_state = await graph.ainvoke(initial_state)
    print("\nИтог: ")
    print(f"  Статус: {final_state['status']}")
    print(f"  Попытки: {final_state['attempts']}")
    print(f"  Результат: {final_state['result']}")

if __name__ == "__main__":
    asyncio.run(main())