#!/usr/bin/env python3
"""
main.py

Самокорректирующийся агент на LangGraph.
- AgentState содержит task, result, attempts, status, error, max_attempts.
- Узлы: execute_task, verify_result, handle_error.
- Инструмент unreliable_tool с ~30% вероятностью бросает ValueError.
- verify_result использует LLM-as-judge (OpenAI Ollama) и отвечает строго "success" или "failed".
- Граф реализует цикл retry до исчерпания max_attempts.
- При запуске выводит номера попыток и финальный статус.
"""

import os
import random
from typing import TypedDict, Optional

from langgraph.graph import StateGraph, START, END
from langchain_openai import ChatOpenAI  # требует OPENAI_API_KEY в окружении
from langchain_core.messages import HumanMessage


# ----------------------------------------------------------------------
# 1. Состояние агента
# ----------------------------------------------------------------------
class AgentState(TypedDict):
    task: str
    result: str
    attempts: int
    status: str          # pending | success | failed | max_attempts
    error: Optional[str]
    max_attempts: int


# ----------------------------------------------------------------------
# 2. Ненадёжный инструмент
# ----------------------------------------------------------------------
def unreliable_tool(task: str) -> str:
    """
    Выполняет простую арифметическую задачу (eval) с вероятностью ~30% бросить ValueError.
    Для демонстрации retry.
    """
    if random.random() < 0.3:
        raise ValueError("Случайная ошибка инструмента")
    # Безопасный eval только для простых арифметических выражений.
    # В реальном коде следует использовать более строгую валидацию.
    return str(eval(task, {"__builtins__": {}}))


# ----------------------------------------------------------------------
# 3. LLM‑judge (OpenAI). Если ключ не задан – fallback на простую проверку.
# ----------------------------------------------------------------------
def get_llm():
    api_key = os.getenv("OPENAI_API_KEY")
    if api_key:
        return ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo")
    return None


LLM = get_llm()


def judge_with_llm(task: str, result: str) -> str:
    """
    Запрашивает у LLM оценку результата.
    Ожидается ответ exatamente "success" или "failed".
    """
    if LLM is None:
        # Fallback: просто сравниваем с правильным ответом через eval
        try:
            correct = str(eval(task, {"__builtins__": {}}))
            return "success" if result.strip() == correct else "failed"
        except Exception:
            return "failed"

    prompt = (
        "You are a judge. Determine if the result correctly answers the task.\n"
        f"Task: {task}\n"
        f"Result: {result}\n"
        "Respond with exactly one word: 'success' if the result is correct, otherwise 'failed'."
    )
    msg = [HumanMessage(content=prompt)]
    response = LLM.invoke(msg)
    verdict = response.content.strip().lower()
    # Приводим к одному из допустимых вариантов
    if "success" in verdict:
        return "success"
    return "failed"


# ----------------------------------------------------------------------
# 4. Узлы графа
# ----------------------------------------------------------------------
def execute_task(state: AgentState) -> AgentState:
    """Выполняет задачу через unreliable_tool."""
    state["attempts"] += 1
    state["error"] = None
    try:
        state["result"] = unreliable_tool(state["task"])
    except Exception as e:
        state["result"] = ""
        state["error"] = str(e)
    return state


def verify_result(state: AgentState) -> AgentState:
    """LLM‑as‑judge: ставит status = success/failed."""
    if state["error"] is not None:
        # Если инструмент уже упал – считаем failed без вызова LLM
        state["status"] = "failed"
        return state

    verdict = judge_with_llm(state["task"], state["result"])
    state["status"] = verdict
    return state


def handle_error(state: AgentState) -> AgentState:
    """Подготовка к повторной попытке."""
    # Ставим статус pending, чтобы граф вернулся к execute_task
    state["status"] = "pending"
    # Ошибку оставляем для логов, но не очищаем – можно очистить, если нужно
    return state


# ----------------------------------------------------------------------
# 5. Построение графа
# ----------------------------------------------------------------------
def build_graph() -> StateGraph:
    workflow = StateGraph(AgentState)

    # Регистрация узлов
    workflow.add_node("execute_task", execute_task)
    workflow.add_node("verify_result", verify_result)
    workflow.add_node("handle_error", handle_error)

    # Входная точка
    workflow.add_edge(START, "execute_task")
    # После выполнения всегда идём к проверке
    workflow.add_edge("execute_task", "verify_result")

    # Условные переходы из verify_result
    def should_continue(state: AgentState) -> str:
        if state["status"] == "success":
            return END
        if state["status"] == "failed" and state["attempts"] < state["max_attempts"]:
            return "handle_error"
        # Либо max_attempts исчерпан, либо иной статус – завершаем
        return END

    workflow.add_conditional_edges(
        "verify_result",
        should_continue,
        {
            "handle_error": "handle_error",
            END: END,
        },
    )

    # После handle_error повторяем выполнение задачи
    workflow.add_edge("handle_error", "execute_task")

    return workflow.compile()


# ----------------------------------------------------------------------
# 6. CLI‑запуск
# ----------------------------------------------------------------------
def main():
    # Пример задачи; можно переопределить через аргумент командной строки
    import sys
    task = sys.argv[1] if len(sys.argv) > 1 else "Вычисли 2+2"

    initial_state: AgentState = {
        "task": task,
        "result": "",
        "attempts": 0,
        "status": "pending",
        "error": None,
        "max_attempts": 5,  # можно изменить
    }

    app = build_graph()
    final_state = app.invoke(initial_state)

    print(f"Задача: {final_state['task']}")
    print(f"Итоговый статус: {final_state['status']}")
    print(f"Попыток сделано: {final_state['attempts']}")
    if final_state["result"]:
        print(f"Результат: {final_state['result']}")
    if final_state["error"]:
        print(f"Последняя ошибка: {final_state['error']}")


if __name__ == "__main__":
    main()