import os
from typing import Annotated
from typing_extensions import TypedDict # Plus robuste pour le State

from dotenv import load_dotenv
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain_google_vertexai import ChatVertexAI

load_dotenv()

# 1. Configuration du modèle (Syntaxe actuelle)
# Assure-hui que GOOGLE_APPLICATION_CREDENTIALS ou GOOGLE_CLOUD_PROJECT est dans ton .env
llm = ChatVertexAI(
    model="gemini-1.5-flash",
    temperature=0,
)

# 2. Définition de l'état
# add_messages permet d'accumuler l'historique au lieu de l'écraser
class State(TypedDict):
    messages: Annotated[list, add_messages]

# 3. Le nœud de l'agent
def call_model(state: State):
    response = llm.invoke(state["messages"])
    # On retourne un dictionnaire qui sera mergé avec l'état actuel
    return {"messages": [response]}

# 4. Construction du graphe (Architecture actuelle)
builder = StateGraph(State)

builder.add_node("agent", call_model)

builder.add_edge(START, "agent")
builder.add_edge("agent", END)

# Compilation
graph = builder.compile()

# 5. Exécution propre
if __name__ == "__main__":
    initial_state = {"messages": [("user", "Salut, tu tournes sur quelle version de Gemini ?")]}
    
    # Utilisation de .stream() pour voir ce qui se passe
    for chunk in graph.stream(initial_state):
        for node, values in chunk.items():
            print(f"--- Node: {node} ---")
            print(values["messages"][-1].content)