from langgraph.graph import StateGraph , START , END 
from langgraph.prebuilt import ToolNode, tools_condition
from Agents import AgentState , agent_analyseur , agent_executor , tools , agent_reporter


# 1. Initialisation du graphe avec l'état personnalisé
workflow = StateGraph(AgentState)

# 2. Ajout des nœuds
workflow.add_node("analyseur", agent_analyseur)
workflow.add_node("executor", agent_executor)
workflow.add_node("reporter", agent_reporter)
workflow.add_node("tools", ToolNode(tools))

# 3. Définition des arêtes stables
workflow.add_edge(START, "analyseur")
workflow.add_edge("analyseur", "executor")
workflow.add_edge("tools", "executor")  
workflow.add_edge("reporter", END)      

# 4. Logique de routage personnalisée
def router(state: AgentState):
    last_message = state["messages"][-1]
    # Si l'exécuteur demande un outil, on y va
    if last_message.tool_calls:
        return "tools"
    # Sinon, cela signifie qu'il a fini son analyse technique -> on passe au rapport
    return "reporter"

# On applique la condition sur l'executor
workflow.add_conditional_edges(
    "executor",
    router,
    {"tools": "tools", "reporter": "reporter"}
)

# 5. Compilation
app = workflow.compile()

# Pour sauvegarder l'image du workflow
with open("graph_workflow.png", "wb") as f:
    f.write(app.get_graph().draw_mermaid_png())

print("Graphique du workflow généré sous : graph_workflow.png")