import os
import asyncio
import random
from typing import TypedDict, Annotated

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 FilesystemBackend, LocalShellBackend, CompositeBackend
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages

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

# ---------- Backend ----------
backend = CompositeBackend([
    LocalShellBackend(workspace_dir="./workspace"),
    FilesystemBackend(),
])

# ---------- Test tool (unreliable) ----------
@tool
def unreliable_tool(query: str) -> str:
    """Tool that succeeds 70% of the time, otherwise raises ValueError."""
    if random.random() < 0.3:
        raise ValueError("Simulated tool failure")
    return f"Result for '{query}'"

# ---------- Agent ----------
agent = create_deep_agent(
    model=llm,
    tools=[unreliable_tool],
    backend=backend,
    system_prompt="You are a helpful agent that can execute tasks and self‑check the result.",
)

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

# ---------- Nodes ----------
async def execute_task(state: AgentState) -> AgentState:
    task = state["task"]
    try:
        # Use the agent to run the task via the tool
        response = await agent.ainvoke(
            {"messages": [HumanMessage(content=task)]},
            {"configurable": {"thread_id": "session-1"}},
        )
        result = response["messages"][-1].content
        state["result"] = result
        state["error"] = None
    except Exception as e:
        state["result"] = ""
        state["error"] = str(e)
    return state

async def verify_result(state: AgentState) -> AgentState:
    # Ask LLM to judge the result
    prompt = (
        f"Task: {state['task']}\n"
        f"Result: {state['result']}\n"
        f"Error: {state['error']}\n"
        "Is the result correct? Respond with only 'success' or 'failed'."
    )
    judge = await llm.ainvoke([HumanMessage(content=prompt)])
    verdict = judge.content.strip().lower()
    if verdict == "success":
        state["status"] = "success"
    else:
        state["status"] = "failed"
    return state

async 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 ----------
graph = StateGraph(AgentState)
graph.add_node("execute_task", execute_task)
graph.add_node("verify_result", verify_result)
graph.add_node("handle_error", handle_error)

graph.set_entry_point("execute_task")
graph.add_conditional_edges(
    "execute_task",
    lambda x: "verify_result",
)
graph.add_conditional_edges(
    "verify_result",
    lambda x: "handle_error" if x["status"] == "failed" else "END",
)
graph.add_conditional_edges(
    "handle_error",
    lambda x: "execute_task" if x["status"] == "pending" else "END",
)

flow = graph.compile()

# ---------- Runner ----------
async def run_task(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 flow.astream(initial_state):
        if event.get("type") == "state":
            state = event["data"]
            print(f"Попытка {state['attempts'] + 1}: status={state['status']}")
            if state["error"]:
                print(f"  Error: {state['error']}")
            if state["result"]:
                print(f"  Result: {state['result']}")
    final_state = event["data"]
    print("\nИтог:", final_state["status"], "за", final_state["attempts"] + 1, "попытки")

if __name__ == "__main__":
    asyncio.run(run_task("Вычисли 2+2", max_attempts=5))