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"
    else:
        return "synthesize"

def build_research_graph():
    """
    Строит граф исследовательского агента
    
    Структура:
    START → outline → research_step → (пока есть пункты) → research_step → synthesize → END
    """
    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_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 = {
        "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"])