import sys
import asyncio
from dotenv import load_dotenv
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_mcp_adapters.client import MultiServerMCPClient

# ✅ LE BON IMPORT POUR LANGGRAPH V1.0+ (et futur v2.0)
from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
from langchain_core.messages import HumanMessage
from langchain_core.messages import ToolMessage

load_dotenv()

async def lancer_agent_moderne():
    llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash")
    
    # Notre configuration MCP habituelle
    client = MultiServerMCPClient({
        "serveur_taches": {
            "command": sys.executable,
            "args": ["serveur_mcp.py"],
            "transport": "stdio"
        }
    })
    outils = await client.get_tools()
    
    async with AsyncSqliteSaver.from_conn_string("memoire_agent.db") as checkpointer:
        
        # 🚀 L'approche moderne :
        # On définit l'agent directement via le constructeur de LangGraph.
        # Plus besoin de 'AgentExecutor', LangGraph gère la boucle lui-même.
        agent = create_react_agent(
            llm, 
            tools=outils, 
            checkpointer=checkpointer,
            interrupt_before=["tools"] # Notre garde-fou
        )
        

        config = {"configurable": {"thread_id": "session_moderne_001"}}
        
        # --- ÉTAPE 1 : L'Agent propose une action ---
        consigne = "Ajoute la tâche 'Apprendre LangGraph' à ma liste."
        print(f"\n🗣️ ORDRE : {consigne}")
        
        async for etape in agent.astream(
            {"messages": [HumanMessage(content=consigne)]}, 
            config=config, 
            stream_mode="values"
        ):
            etape["messages"][-1].pretty_print()

        # --- ÉTAPE 2 : Vérification de l'état ---
        etat_actuel = await agent.aget_state(config)
        
        if etat_actuel.next:
            print(f"\n⚠️ ACTION EN ATTENTE : L'agent veut utiliser l'outil : {etat_actuel.next}")
            confirmation = input("👉 Voulez-vous autoriser cette action ? (oui/non) : ")
            
            if confirmation.lower() == "oui":
                # --- ÉTAPE 3 : On relance l'agent là où il s'est arrêté ---
                print("🚀 Validation reçue. Exécution en cours...")
                async for etape in agent.astream(None, config=config, stream_mode="values"):
                    etape["messages"][-1].pretty_print()
            else:
                print("❌ Action annulée par l'utilisateur.")
                # On crée un message de résultat d'outil manuel pour clore la boucle
                outil_id = etape["messages"][-1].tool_calls[0]["id"]
                message_annulation = ToolMessage(
                    tool_call_id=outil_id,
                    content="L'utilisateur a refusé cette action. Propose une alternative ou arrête-toi là."
                )
                
                # On met à jour l'état de l'agent avec ce message de "faux" résultat
                await agent.aupdate_state(config, {"messages": [message_annulation]})
                
                # On relance pour que l'IA réagisse à l'annulation
                async for etape in agent.astream(None, config=config, stream_mode="values"):
                    etape["messages"][-1].pretty_print()

        else:
            print("\n✅ L'agent a terminé sans avoir besoin de validation.")

if __name__ == "__main__":
    asyncio.run(lancer_agent_moderne())