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

def should_continue(state: ReflectState) -> str:
    """Функция-условие: продолжать или завершить"""
    # Проверяем вердикт
    if state["verdict"] == "ok":
        print(f"✅ Ответ принят после {state['round']} раундов")
        return "end"
    
    # Проверяем лимит раундов
    if state["round"] >= state["max_rounds"]:
        print(f"⚠️ Достигнут лимит раундов ({state['max_rounds']}). Возвращаем последнюю версию.")
        return "end"
    
    # Требуется доработка
    print(f"🔄 Запускаем раунд доработки {state['round'] + 1}/{state['max_rounds']}")
    return "rewrite"

def build_reflection_graph(max_rounds: int = 2):
    """Строит граф с рефлексией"""
    # Инициализация графа с состоянием
    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,
        {
            "rewrite": "rewrite",
            "end": END
        }
    )
    
    # После rewrite возвращаемся на reflect для повторной оценки
    graph.add_edge("rewrite", "reflect")
    
    # Компилируем граф
    return graph.compile()

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