from langgraph.graph import StateGraph, START, END
from state import BriefState
from nodes import outline_node, research_step_node, synthesize_node

def should_continue_research(state: BriefState) -> str:
    """Условие: продолжать исследование или переходить к синтезу"""
    if state["step_index"] < len(state["outline"]):
        return "continue_research"
    else:
        return "synthesize"

def build_research_graph():
    """Строит граф исследовательского агента"""
    # Инициализация графа с состоянием
    graph = StateGraph(BriefState)
    
    # Добавляем узлы
    graph.add_node("outline", outline_node)
    graph.add_node("research_step", research_step_node)
    graph.add_node("synthesize", synthesize_node)
    
    # Добавляем рёбра
    graph.add_edge(START, "outline")
    graph.add_edge("outline", "research_step")
    
    # Условное ребро после research_step
    graph.add_conditional_edges(
        "research_step",
        should_continue_research,
        {
            "continue_research": "research_step",  # Цикл по шагам
            "synthesize": "synthesize"
        }
    )
    
    # После синтеза завершаем
    graph.add_edge("synthesize", END)
    
    # Компилируем граф
    return graph.compile()

# Пример использования (для тестирования)
if __name__ == "__main__":
    from dotenv import load_dotenv
    load_dotenv()
    
    # Создаём граф
    app = build_research_graph()
    
    # Входное состояние
    initial_state: BriefState = {
        "topic": "Как студенту безопасно подключать MCP к LangChain",
        "outline": None,
        "step_index": 0,
        "notes": [],
        "final_brief": None
    }
    
    # Запуск
    final_state = app.invoke(initial_state)
    
    # Вывод результата
    print("\n" + "="*80)
    print("ИТОГОВЫЙ ИССЛЕДОВАТЕЛЬСКИЙ БРИФ")
    print("="*80)
    print(final_state["final_brief"])