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
from langchain_core.messages import AIMessage , ToolMessage


# 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("reporter", END)      

# 4. Logique de routage personnalisée
def router(state: AgentState):
    messages = state["messages"]
    last_message = messages[-1]

    # CAS 1 : L'agent demande un outil
    if hasattr(last_message, "tool_calls") and last_message.tool_calls:
        return "tools"

    # CAS 2 : Analyse du retour de l'outil 
    if isinstance(last_message, ToolMessage):
        content_upper = last_message.content.upper()
        
        # On détecte ton marqueur spécifique ou le mot "ERROR"
        if "ERREUR PYTHON" in content_upper or "ERROR" in content_upper:
            print("--- LOG : Erreur détectée, renvoi à l'exécuteur pour correction ---")
            return "executor" 
        
        # Si pas d'erreur, on peut passer à la synthèse
        return "reporter"

    # CAS 3 : Par défaut
    return "reporter"

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

workflow.add_conditional_edges(
    "tools" , 
    router , 
    {"executor": "executor", "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")