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

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

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

import random
from typing import TypedDict, Dict, Any

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. Инструмент – unreliable_tool
# --------------------------------------------------------------------------- #

def unreliable_tool(input_text: str) -> str:
    """
    Случайно бросает ValueError с вероятностью ~30%.
    Иначе возвращает результат выполнения простого арифметического выражения.
    """
    if random.random() < 0.3:
        raise ValueError("Инструмент случайно упал")
    # Простейший парсер: ищем числа и оператор +, -, *, /
    try:
        left, op, right = input_text.split()
        left_val = float(left)
        right_val = float(right)
        if op == '+':
            return str(left_val + right_val)
        elif op == '-':
            return str(left_val - right_val)
        elif op == '*':
            return str(left_val * right_val)
        elif op == '/':
            return str(left_val / right_val)
    except Exception:
        pass
    # Если не удалось распарсить – просто возвращаем строку
    return f"результат: {input_text}"


# --------------------------------------------------------------------------- #
# 3. Узлы графа
# --------------------------------------------------------------------------- #

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


def verify_result(state: AgentState) -> Dict[str, Any]:
    """
    LLM‑оценка результата. Модель должна ответить «success» или «failed».
    Если результат пустой – сразу считается «failed».
    """
    if not state["result"]:
        return {"status": "failed"}

    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
    prompt = (
        f"Задача: {state['task']}\n"
        f"Результат: {state['result']}\n\n"
        "Оцените результат. Ответьте только одним словом: success или failed."
    )
    response = llm.invoke(prompt).content.strip().lower()
    if response.startswith("success"):
        return {"status": "success"}
    else:
        return {"status": "failed"}


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


# --------------------------------------------------------------------------- #
# 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_conditional_edges(
    "execute_task",
    lambda x: x["status"],
    {
        "pending": "verify_result",
        "failed": "handle_error",
    },
)

builder.add_conditional_edges(
    "verify_result",
    lambda x: x["status"],
    {
        "success": END,
        "failed": "handle_error",
    },
)

builder.add_conditional_edges(
    "handle_error",
    lambda x: x["status"],
    {
        "pending": "execute_task",
        "max_attempts": END,
    },
)

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


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

def main() -> None:
    """
    CLI‑интерфейс: пользователь вводит задачу, например «2 + 2».
    Выводятся номера попыток и финальный статус.
    """
    task = input("Введите задачу (пример: 2 + 2): ").strip()
    if not task:
        print("Задача не может быть пустой.")
        return

    # Инициализируем состояние
    initial_state: AgentState = {
        "task": task,
        "result": "",
        "attempts": 0,
        "status": "pending",
        "error": None,
        "max_attempts": 5,
    }

    print("\nЗапуск агента...\n")

    # Запускаем граф
    for step in graph.stream(initial_state, stream_mode="values"):
        state = step.values
        attempts = state["attempts"] + 1 if state["status"] != "pending" else state["attempts"]
        print(f"Попытка {attempts}:")
        if state.get("error"):
            print(f"  Ошибка инструмента: {state['error']}")
        else:
            print(f"  Результат: {state['result']}")

        if state["status"] == "success":
            print("\nИтог: success за", attempts, "попыток")
            break
        elif state["status"] == "max_attempts":
            print("\nИтог: max_attempts – задача не решена после 5 попыток")
            break

    else:
        # Если цикл завершился без явного выхода (должно быть)
        print("\nНеизвестный конец выполнения.")


if __name__ == "__main__":
    main()