from langgraph.graph import StateGraph, END
from app.agent.state import AgentState
from app.agent.nodes import (
    generate_initial_questions,
    process_answers,
    analyze_round,
    generate_checklist,
    check_round_complete
)


def create_checklist_agent() -> StateGraph:
    """Создает LangGraph для чеклист-агента"""
    
    # Создаем граф с состоянием AgentState
    workflow = StateGraph(AgentState)
    
    # Добавляем ноды
    workflow.add_node("generate_initial_questions", generate_initial_questions)
    workflow.add_node("process_answers", process_answers)
    workflow.add_node("analyze_round", analyze_round)
    workflow.add_node("generate_checklist", generate_checklist)
    
    # Устанавливаем начальную точку
    workflow.set_entry_point("generate_initial_questions")
    
    # Добавляем переходы
    # После генерации вопросов - ждем ответы (END чтобы вернуть контроль)
    workflow.add_edge("generate_initial_questions", END)
    
    # После обработки ответов - анализируем раунд
    workflow.add_edge("process_answers", "analyze_round")
    
    # После анализа - либо ждем новые ответы, либо генерируем чеклист
    workflow.add_conditional_edges(
        "analyze_round",
        check_round_complete,
        {
            "wait_for_answers": END,  # Ждем следующие ответы
            "generate_checklist": "generate_checklist"  # Генерируем чеклист
        }
    )
    
    # После генерации чеклиста - конец
    workflow.add_edge("generate_checklist", END)
    
    return workflow.compile()


# Создаем экземпляр агента
checklist_agent = create_checklist_agent()