from langgraph.graph import START, END
from langgraph.graph.state import CompiledStateGraph
from langgraph.checkpoint.memory import InMemorySaver

from utils.nodes import *
from utils.state import getState
from utils.tools import getTools

def getGraph()->CompiledStateGraph:
    """
    Récupérer le graphe compilé et prêt à invoquer

    Returns:
        CompiledStateGraph: Graphe compilé
    """
    workflow = getState() # State prêt à utiliser

    # Définition des sommets du graphe
    workflow.add_node(user_prompt)
    workflow.add_node(LLM_central)
    workflow.add_node(preparation_docs)
    workflow.add_node(inject_preparation_prompt)
    workflow.add_node("tool_node", tool_node)# BasicToolNode(tools=getTools())) # N'est pas une fonction, mais une classe instanciée, je dois précisier le nom du node
    workflow.add_node("weekly_report_tools", weekly_report_tools)
    workflow.add_node(context_shortener) # Réduit la taille du contexte
    workflow.add_node("context_shortener_2", context_shortener) # Le même, sous un autre nom pour le différencier dans le graphe

    # Arrêtes
    workflow.set_conditional_entry_point(is_resumes_reports_already_initialised, {
        "résumés non disponibles": "inject_preparation_prompt", # Résumés non générés
        "résumés déjà générés": "user_prompt" # Résumés déjà prêts, je peux aller direct à la partie principale
    })
    workflow.add_edge("inject_preparation_prompt", "preparation_docs")
    workflow.add_conditional_edges("preparation_docs", should_continue, {
        "tools":"weekly_report_tools",
        "no_tools":"context_shortener" # FIN de la préparation, on réduit le contexte avant de passer à la suite
    })
    workflow.add_edge("context_shortener", "user_prompt") # Et ici, je rejoins la partie principale qui rédigera le rapport
    workflow.add_conditional_edges("user_prompt", lambda state: END if state['stop'] else "continue", {
        END: END,
        "continue": "LLM_central"
    })

    workflow.add_edge("weekly_report_tools", "preparation_docs")
    workflow.add_conditional_edges("tool_node", should_shorten, {
        'sous la limite': "LLM_central",
        'réduire contexte': "context_shortener_2"
    })
    workflow.add_edge("context_shortener_2", "LLM_central")
    workflow.add_conditional_edges("LLM_central", should_continue, {
        "tools":"tool_node",
        "no_tools":"user_prompt"
    })

    return workflow.compile(checkpointer=InMemorySaver()) # TODO: Rempalcer par une vrai BDD de prod

if __name__ == "__main__":
    # Affichage du graphe
    getGraph().get_graph().draw_mermaid_png(output_file_path="imgs/agent.png")