from langgraph.graph import StateGraph, START, END

# ⚠️ Import corrigé (selon version)
try:
    from langgraph.prebuilt import ToolNode
except ImportError:
    # fallback si ToolNode n'existe pas
    from langchain_core.runnables import RunnableLambda

    def ToolNode(tools):
        def run_tools(state):
            return state  # à adapter si besoin
        return RunnableLambda(run_tools)

# Assure-toi d'importer également agent_llm_vision depuis ton fichier Agents
from Agents import AgentState, agent_extracteur, agent_builder, agent_ocr
# 💡 Note : Pense à ajouter la fonction agent_llm_vision dans ton fichier Agents.py !
try:
    from Agents import agent_llm_vision
except ImportError:
    # Fallback temporaire si tu ne l'as pas encore écrit dans Agents.py
    def agent_llm_vision(state: AgentState):
        print("🤖 Exécution de l'Agent LLM Vision...")
        return {"messages": ["Traitement vision effectué"]}

workflow = StateGraph(AgentState)

# 1. Définition des Nœuds
workflow.add_node("agent_ocr", agent_ocr)
workflow.add_node("agent_llm_vision", agent_llm_vision) # <-- Nouveau nœud rouge
workflow.add_node("agent_extracteur", agent_extracteur)
workflow.add_node("agent_builder", agent_builder)


# 2. Fonction de routage conditionnel
def route_after_ocr(state: AgentState):
    """
    Cette fonction regarde le flag 'use_vision' retourné par agent_ocr
    pour choisir le chemin dans le graphe.
    """
    if state.get("use_vision") is True:
        return "vision_path"
    else:
        return "extracteur_path"

# 3. Définition des Arêtes (Edges)
workflow.add_edge(START, "agent_ocr")

# 🔀 Flèche conditionnelle (La logique "If true" de ton schéma)
workflow.add_conditional_edges(
    "agent_ocr",
    route_after_ocr,
    {
        "vision_path": "agent_llm_vision",     # Si use_vision == True -> Va vers LLM Vision
        "extracteur_path": "agent_extracteur"  # Si use_vision == False -> Va vers agent_extracteur
    }
)

# Suite et fin des flux
workflow.add_edge("agent_extracteur", "agent_builder")
workflow.add_edge("agent_llm_vision", "agent_builder")
workflow.add_edge("agent_builder", END)




# 5. Compilation
app = workflow.compile() 

# 6. Graph
try:
    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")
except Exception as e:
    print(f"  Erreur génération image : {e}")