import os
from dotenv import load_dotenv
from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings
from langchain_qdrant import QdrantVectorStore
from qdrant_client import QdrantClient
from langchain_classic.chains import create_retrieval_chain
from langchain_classic.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate

# 1. Chargement de la clé API
load_dotenv()
api_key = os.getenv("GOOGLE_API_KEY")

def interroger_mon_rag(question):
    # --- A. CONFIGURATION DES MODÈLES ---
    # Le modèle qui va générer la réponse
    llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", google_api_key=api_key, temperature=0)
    
    # Le modèle qui transforme le texte en vecteurs (doit être le même que celui utilisé pour remplir la base)
    embeddings = GoogleGenerativeAIEmbeddings(model="models/gemini-embedding-001")

    # --- B. CONNEXION À QDRANT ---
    # Remplacez l'URL par l'adresse de votre VM (ex: http://localhost:6333)
    # Et mettez le vrai nom de votre collection Qdrant
    qdrant_url = "http://qdrant:6333" 
    nom_collection = "rag-1"

    client = QdrantClient(url=qdrant_url)
    
    # On connecte LangChain à Qdrant
    vector_store = QdrantVectorStore(
        client=client,
        collection_name=nom_collection,
        embedding=embeddings
    )
    
    # On transforme la base en "Retriever" (un outil qui va chercher les 3 meilleurs documents)
    retriever = vector_store.as_retriever(search_kwargs={"k": 3})

    # --- C. CRÉATION DU PROMPT RAG ---
    # Le mot {context} sera automatiquement remplacé par les documents trouvés dans Qdrant
    system_prompt = (
        "Tu es un assistant utile. Utilise UNIQUEMENT les morceaux de contexte "
        "fournis ci-dessous pour répondre à la question. Si tu ne connais pas la réponse "
        "à partir du contexte, dis simplement que tu ne sais pas."
        "\n\n"
        "Contexte :\n{context}"
    )
    
    prompt = ChatPromptTemplate.from_messages([
        ("system", system_prompt),
        ("human", "{input}"),
    ])

    # --- D. ASSEMBLAGE DE LA CHAÎNE ---
    # 1. Chaîne qui combine les documents récupérés et le prompt
    question_answer_chain = create_stuff_documents_chain(llm, prompt)
    
    # 2. Chaîne globale qui fait la recherche PUIS appelle la chaîne précédente
    rag_chain = create_retrieval_chain(retriever, question_answer_chain)

    # --- E. EXÉCUTION ---
    print(f"Recherche de la réponse pour : '{question}'...\n")
    reponse = rag_chain.invoke({"input": question})
    
    print("🤖 Réponse de Gemini :")
    print(reponse["answer"])
    
    print("\n📄 Documents sources utilisés :")
    for doc in reponse["context"]:
        # Affiche un extrait du document source et ses métadonnées
        print(f"- {doc.page_content[:100]}... (Source: {doc.metadata})")

# Test du script
if __name__ == "__main__":
    interroger_mon_rag("Quelles sont les passions de vincent poupet ?")