#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
main.py – запуск LangGraph‑агента с самопроверкой и повторными попытками.

Требования:
    - Python 3.10+
    - pip install langgraph langchain-openai
"""

import random
from typing import TypedDict, Dict

from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import InMemorySaver


# --------------------------------------------------------------------------- #
# 1. Состояние графа
# --------------------------------------------------------------------------- #

class AgentState(TypedDict):
    task: str          # исходная задача
    result: str        # результат выполнения инструмента
    attempts: int      # число попыток
    status: str        # pending | success | failed | max_attempts
    error: str | None  # сообщение об ошибке, если есть
    max_attempts: int  # лимит повторов


# --------------------------------------------------------------------------- #
# 2. Инструмент – «не надёжный» калькулятор
# --------------------------------------------------------------------------- #

def unreliable_tool(input_str: str) -> str:
    """
    Случайно бросает ValueError с вероятностью ~30%.
    Иначе возвращает результат простого арифметического выражения.
    """
    if random.random() < 0.3:
        raise ValueError("Инструмент временно недоступен")
    try:
        # eval безопасный для простых выражений (здесь только числа и +)
        return str(eval(input_str))
    except Exception as exc:
        raise ValueError(f"Невозможно вычислить: {exc}") from exc


# --------------------------------------------------------------------------- #
# 3. Узлы
# --------------------------------------------------------------------------- #

def execute_task(state: AgentState) -> Dict[str, object]:
    """
    Выполняет задачу через unreliable_tool.
    При ошибке сохраняет сообщение об ошибке и помечает статус как failed.
    """
    try:
        result = unreliable_tool(state["task"])
        return {
            "result": result,
            "error": None,
            "status": "pending",  # результат получен, но ещё не проверён
        }
    except Exception as exc:
        return {
            "result": "",
            "error": str(exc),
            "status": "failed",
        }


def verify_result(state: AgentState) -> Dict[str, object]:
    """
    LLM‑проверка результата. Модель должна вернуть ровно 'success' или 'failed'.
    """
    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
    prompt = (
        f"Задача: {state['task']}\n"
        f"Результат: {state['result']}\n"
        "Оцените результат. Ответьте только одним словом: 'success' или 'failed'."
    )
    response = llm.invoke(prompt).content.strip().lower()
    if response not in {"success", "failed"}:
        # Если модель дала непонятный ответ – считаем это ошибкой
        return {
            "status": "failed",
            "error": f"LLM дал неожиданный ответ: {response}",
        }
    return {"status": response}


def handle_error(state: AgentState) -> Dict[str, object]:
    """
    Увеличиваем счётчик попыток и готовим к повтору.
    Если лимит превышен – помечаем как max_attempts.
    """
    attempts = state["attempts"] + 1
    if attempts >= state["max_attempts"]:
        return {"status": "max_attempts", "attempts": attempts}
    return {"status": "pending", "attempts": attempts}


# --------------------------------------------------------------------------- #
# 4. Создание графа
# --------------------------------------------------------------------------- #

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")

# Сохраняем состояние в памяти (для отладки)
memory = InMemorySaver()
graph = builder.compile(persist_to_db=memory)


# --------------------------------------------------------------------------- #
# 5. Запуск
# --------------------------------------------------------------------------- #

def main() -> None:
    # Пример задачи: простое арифметическое выражение
    task_text = "2 + 2"

    initial_state: AgentState = {
        "task": task_text,
        "result": "",
        "attempts": 0,
        "status": "pending",
        "error": None,
        "max_attempts": 5,
    }

    # Запускаем граф
    final_state = graph.invoke(initial_state)

    # Выводим результаты
    print(f"Задача: {task_text}")
    for i in range(final_state["attempts"] + 1):
        state_snapshot = memory.get(i)
        if state_snapshot is None:
            continue
        attempt_num = state_snapshot["attempts"]
        status = state_snapshot["status"]
        error_msg = state_snapshot.get("error")
        result = state_snapshot.get("result", "")
        print(
            f"Попытка {attempt_num + 1}: "
            f"{'Error: ' + error_msg if error_msg else 'Результат: ' + result} → "
            f"verify: {status}"
        )
    print(f"\nИтог: {final_state['status']} за {final_state['attempts'] + 1} попыток")


if __name__ == "__main__":
    main()