from langgraph.graph import StateGraph, START, END
from state import ReflectState
from nodes import draft_answer, reflect, rewrite

def should_continue_or_retry(state: ReflectState) -> str:
    """
    Расширенная логика: сначала проверяем try/except ошибки,
    затем вердикт критика
    """
    # 1. Проверяем ошибки от try/except в draft_answer
    if state.get("error"):
        print(f"   ⚠️ Обнаружена ошибка генерации: {state['error']}")
        
        if state["round"] < state["max_rounds"]:
            print(f"   🔄 Повторная попытка генерации ({state['round'] + 1}/{state['max_rounds']})...")
            return "retry_draft"
        else:
            print(f"   ❌ Достигнут лимит попыток, возвращаем последнюю версию")
            return "end"
    
    # 2. Проверяем вердикт критика
    if state["verdict"] == "ok":
        print(f"✅ Ответ принят после {state['round']} раундов")
        return "end"
    
    # 3. Проверяем лимит раундов доработки
    if state["round"] >= state["max_rounds"]:
        print(f"⚠️ Достигнут лимит раундов ({state['max_rounds']}). Возвращаем последнюю версию.")
        return "end"
    
    # 4. Требуется доработка
    print(f"🔄 Запускаем раунд доработки {state['round'] + 1}/{state['max_rounds']}")
    return "rewrite"

def build_reflection_graph(max_rounds: int = 2, max_retries: int = 3):
    """
    Строит граф с рефлексией и retry при ошибках
    
    Структура:
    START → draft_answer → reflect
             ↑              ↓
             |         needs_revision & round < max_rounds → rewrite → reflect
             |              ↓
             |              ok → END
             |              ↓
             |         error & round < max_retries → draft_answer (retry)
             |              ↓
             └──────────────┘
    """
    graph = StateGraph(ReflectState)
    
    # Добавляем узлы
    graph.add_node("draft_answer", draft_answer)
    graph.add_node("reflect", reflect)
    graph.add_node("rewrite", rewrite)
    
    # Добавляем рёбра
    graph.add_edge(START, "draft_answer")
    graph.add_edge("draft_answer", "reflect")
    
    # Условное ребро после рефлексии (или после ошибки)
    graph.add_conditional_edges(
        "reflect",
        should_continue_or_retry,
        {
            "retry_draft": "draft_answer",  # retry при ошибке генерации
            "rewrite": "rewrite",            # доработка по вердикту критика
            "end": END
        }
    )
    
    # После rewrite возвращаемся на reflect для повторной оценки
    graph.add_edge("rewrite", "reflect")
    
    # Компилируем граф
    return graph.compile()

# Пример использования
if __name__ == "__main__":
    # Создаём граф
    app = build_reflection_graph(max_rounds=2, max_retries=3)
    
    # Входное состояние
    initial_state: ReflectState = {
        "question": "Объясни студенту разницу между tool и resource в MCP",
        "draft": "",
        "critique": "",
        "verdict": "needs_revision",
        "round": 1,
        "max_rounds": 2,
        "error": None
    }
    
    # Запуск
    final_state = app.invoke(initial_state)
    
    # Вывод результата
    print("\n" + "="*80)
    print("ФИНАЛЬНЫЙ ОТВЕТ:")
    print("="*80)
    print(final_state["draft"])