from langgraph.graph import StateGraph, START, END
from state import ReflectState
from nodes import generate_answer

def should_retry(state: ReflectState) -> str:
    """
    Условие: продолжать попытки или завершить
    Проверяет наличие ошибки и количество попыток
    """
    # Если ошибки нет - успех
    if state["error"] is None:
        print(f"✅ Ответ успешно сгенерирован за {state['round']} попыт(ок/ки)")
        return "end"
    
    # Если есть ошибка и не превышен лимит - повторяем
    if state["round"] < state["max_rounds"]:
        print(f"🔄 Повторная попытка ({state['round'] + 1}/{state['max_rounds']})...")
        return "retry"
    
    # Превышен лимит - завершаем с ошибкой
    print(f"⚠️ Достигнут лимит попыток ({state['max_rounds']}). Последняя ошибка: {state['error']}")
    return "end"

def build_retry_graph(max_rounds: int = 3):
    """
    Строит граф с retry-логикой через try/except
    
    Структура:
    START → generate_answer → should_retry
        → (если error None) → END
        → (если error и round < max_rounds) → generate_answer (с увеличенным round)
        → иначе → END
    """
    graph = StateGraph(ReflectState)
    
    # Добавляем единственный узел (он же будет вызываться повторно)
    graph.add_node("generate_answer", generate_answer)
    
    # Добавляем рёбра
    graph.add_edge(START, "generate_answer")
    
    # Условное ребро после генерации
    graph.add_conditional_edges(
        "generate_answer",
        should_retry,
        {
            "retry": "generate_answer",  # Петля для повторной попытки
            "end": END
        }
    )
    
    return graph.compile()

# Пример использования
if __name__ == "__main__":
    # Создаём граф
    app = build_retry_graph(max_rounds=3)
    
    # Входное состояние
    initial_state = {
        "question": "Объясни студенту разницу между tool и resource в MCP",
        "draft": "",
        "error": None,
        "round": 1,
        "max_rounds": 3
    }
    
    # Запуск
    final_state = app.invoke(initial_state)
    
    # Вывод результата
    print("\n" + "="*80)
    if final_state["error"] is None:
        print("✅ ФИНАЛЬНЫЙ ОТВЕТ:")
        print("="*80)
        print(final_state["draft"])
    else:
        print("❌ НЕ УДАЛОСЬ СГЕНЕРИРОВАТЬ КАЧЕСТВЕННЫЙ ОТВЕТ")
        print(f"Ошибка: {final_state['error']}")