import pinecone
import google.generativeai as genai
from sklearn.metrics.pairwise import cosine_similarity
from sentence_transformers import SentenceTransformer
import numpy as np
import gradio as gr
from pinecone import Pinecone

# Configurer l'API Google
genai.configure(api_key="AIzaSyBLlaGtxtzHmVgMfOC02AfgvOoKTwXGGIc")
pc = Pinecone(api_key="pcsk_2vJdv7_KGHuiapnC85fgFaUGdZKyqiHPBv4hPMSnDq7FnzLSvhs2Ur51C3Z6rmbyyYVLZ5", environment="us-west1-gcp")

# 2. Initialisation
index_name = "index-embeddings-code-du-travail232333"
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
model_gene = genai.GenerativeModel("gemini-2.0-flash")

# Connecter à l'index Pinecone
index = pc.Index(index_name)

# Fonction pour récupérer les 3 résultats les plus similaires de Pinecone
def get_most_similar_chunk(query):
    # Utiliser le modèle SentenceTransformer pour générer l'embedding de la requête
    model = SentenceTransformer("all-MiniLM-L6-v2")
    query_embedding = model.encode(query)

    # Effectuer une recherche dans Pinecone
    result = index.query(
        vector=query_embedding.tolist(),  # Convertir en liste pour Pinecone
        top_k=5,  # Nombre de résultats à retourner
        include_metadata=True  # Inclure les métadonnées
    )

    # Vérifier si des matches sont trouvés
    if 'matches' not in result:
        return []

    # Récupérer les chunks et leurs scores de similarité
    similar_chunks = []
    for match in result['matches']:
        if 'metadata' in match:
            # Vérifier que la clé 'text' existe dans les métadonnées
            text = match['metadata'].get('text', 'Texte non disponible')
            similar_chunks.append((text, match['score']))  # Utiliser 'text' ici, pas 'chunk'

    return similar_chunks

# Initialisation du modèle de génération
model_gene = genai.GenerativeModel("gemini-2.0-flash")

def chatbot(user_prompt):
    if not user_prompt:
        return "❗ Veuillez entrer une question."

    similar_chunks = get_most_similar_chunk(user_prompt)
    context = "\n".join(
        [f"Contexte {i+1}: {chunk} (score: {score:.2f})" for i, (chunk, score) in enumerate(similar_chunks)]
    )

    prompt = f"""
        TTu es un assistant juridique expert en droit du travail, et ta mission est d’apporter des réponses claires, concises et accessibles à toutes les 
        questions liées au Code du travail, en te basant sur les informations disponibles dans le documents du code du travail fournit et ajoute des réference à partir du pdf. 
        Lorsque quelqu’un te pose une question, veille à ce que ta réponse soit facile à comprendre, sans jargon complexe. Si tu rencontres un manque 
        d’informations pour répondre de manière précise, sois transparent à ce sujet et donne des exemples concrets pour illustrer la situation. 
        Chaque réponse doit commencer par une accroche sympathique pour engager l’utilisateur, puis expliquer de manière simple les concepts clés 
        en lien avec la question posée. Enfin, offre une mini check-list ou résumé des points essentiels pour aider l’utilisateur à mieux comprendre.
        Question : {user_prompt}\n
        {context}\nRéponse :
    """

    response = model_gene.generate_content(prompt)
    return response.text

# 5. Interface Gradio
demo = gr.Interface(
    fn=chatbot,
    inputs=gr.Textbox(label="💬 Entrez votre question ici :"),
    outputs=gr.Textbox(label="📢 Réponse du chatbot :"),
    title="🤖 Assistant Code du travail",
    description="Posez votre question sur le Code du travail, l’assistant vous répond avec des sources pertinentes."
)

if __name__ == "__main__":
    demo.launch()