import google.generativeai as genai
import gradio as gr
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
from datasets import load_dataset
import pandas as pd

# Charger les embeddings sauvegardés
dataset = load_dataset("ayouubelb/embeddings")
data = pd.DataFrame(dataset)
embeddings = np.array(data['embedding']).astype('float32')
texts = data['metadata']

# Créer un index FAISS
dim = embeddings.shape[1]
index = faiss.IndexFlatIP(dim) 
index.add(embeddings)

genai.configure(api_key="AIzaSyBLlaGtxtzHmVgMfOC02AfgvOoKTwXGGIc")
model_gene = genai.GenerativeModel("gemini-2.0-flash")
model =  SentenceTransformer("dangvantuan/sentence-camembert-base")

def rag_legal_bot(user_input):
    # Générer l'embedding de la requête
    query_embedding = model.encode([user_input]).astype("float32")
    distances, indices = index.search(query_embedding, 5)

    print(distances, indices)
    # Contexte
    contextes = [texts[i] for i in indices[0]]
    contexte = "\n\n".join([f"[Contexte {i+1}]\n{ctx}" for i, ctx in enumerate(contextes)])
    
    #le prompt pour Gemini
    prompt = f"""
    Tu es un assistant juridique expert en droit français. Ta mission est de fournir des réponses précises, 
    factuelles, claires et pédagogiques, destinées à un public mixte (non-juristes et professionnels). 
    Tu t’appuies uniquement sur les textes de loi et contextes fournis, sans interpréter ni spéculer.

    Consignes détaillées :
    
    1. Rigueur juridique :
    - Cite précisément les articles de loi pertinents.
    - Mentionne toujours le critère ou principe juridique central lié à la question (ex. lien de subordination pour contrat de travail).
    - Ne réponds qu’à partir du contexte et des lois fournies.
    
    2. Structure et pédagogie :
    - Organise ta réponse avec titres, sous-titres et listes à puces pour faciliter la lecture.
    - Explique clairement les notions juridiques, en évitant le jargon technique inutile.
    - Utilise un exemple concret chaque fois que possible pour illustrer les concepts.
    
    3. Portée juridique :
    - Indique la portée, les conditions d’application et les conséquences juridiques des règles.
    - Mentionne les enjeux pratiques ou contentieux liés au sujet (ex. risques de requalification, effets sur les droits).
    - Précise quand des informations complémentaires seraient nécessaires.
    
    4. Objectivité et professionnalisme :
    - Reste neutre, rigoureux et factuel, sans jugement ni supposition.
    - Si un élément manque ou n’est pas clair, indique-le explicitement.
    - Termine toujours par un résumé clair et synthétique des points essentiels.
    
    Exemple d’introduction dans la réponse :
    « Selon l’article X du Code Y, … »
    « Le critère principal pour distinguer est … »
    « Cette règle a pour conséquence … »
    
    Ce prompt incite à :
    
    - Mettre en avant le critère clé dans la réponse (ex. lien de subordination)
    - Donner la portée pratique et juridique des règles
    - Structurer la réponse de façon pédagogique et professionnelle
    - Signaler les limites éventuelles du contexte fourni

    Question : {user_input}

    {contexte}

    Réponse :
    """

    response = model_gene.generate_content(prompt)
    print("Contexte utilisé :\n", contextes)
    return response.text



#Interface Gradio
interface = gr.Interface(
    fn=rag_legal_bot,
    inputs=gr.Textbox(label="💬 Entrez votre question ici :"),
    outputs=gr.Textbox(label="📢 Réponse du chatbot :"),
    title="🧑‍⚖️ Assistant Juridique",
    description="Posez votre question juridique."
)

interface.launch()