import os
from huggingface_hub import login
from datasets import load_dataset
import gradio as gr
from llama_cpp import Llama
from huggingface_hub import hf_hub_download
import chromadb
from sentence_transformers import SentenceTransformer
import time

# Charger le token depuis les secrets
hf_token = os.getenv("HF_TOKEN")  # Assurez-vous que 'HF_TOKEN' est bien le nom du secret Hugging Face

# Connecte-toi à Hugging Face
login(hf_token)

# Charger le dataset une seule fois
dataset = load_dataset("Maryem2025/dataset-test")  # Chargez le dataset une fois

# Initialisation du modèle Llama avec une taille de contexte réduite
llm = Llama(
    model_path=hf_hub_download(
        repo_id="TheBloke/CapybaraHermes-2.5-Mistral-7B-GGUF",
        filename="capybarahermes-2.5-mistral-7b.Q2_K.gguf",
    ),
    n_ctx=1024,  # Réduire la taille du contexte
    n_gpu_layers=50,  # Ajustez selon votre VRAM
)

# Initialisation de ChromaDB Vector Store
class VectorStore:
    def __init__(self, collection_name, batch_size=10):
        self.embedding_model = SentenceTransformer('sentence-transformers/multi-qa-MiniLM-L6-cos-v1')
        self.chroma_client = chromadb.Client()
        self.batch_size = batch_size

        # Supprimer la collection existante si elle existe
        if collection_name in self.chroma_client.list_collections():
            self.chroma_client.delete_collection(collection_name)
        
        # Créer une nouvelle collection
        self.collection = self.chroma_client.create_collection(name=collection_name)

    def populate_vectors(self, dataset):
        # Sélectionner les colonnes pertinentes à concaténer
        names = dataset['train']['name'][:200]
        ingredients = dataset['train']['ingredients'][:200]
        instructions = dataset['train']['instructions'][:200]
        cuisine = dataset['train']['cuisine'][:200]
        total_time = dataset['train']['total_time'][:200]

        # Concaténer les textes à partir des colonnes sélectionnées
        texts = [
            f"Name: {name}. Ingredients: {ingr}. Instructions: {instr}. Cuisine: {cui}. Total time: {total} minutes."
            for name, ingr, instr, cui, total in zip(names, ingredients, instructions, cuisine, total_time)
        ]

        embeddings_batch = []
        documents_batch = []

        for i, item in enumerate(texts):
            embeddings = self.embedding_model.encode(item).tolist()
            embeddings_batch.append(embeddings)
            documents_batch.append(item)

            # Quand le batch est plein, on ajoute les embeddings
            if len(embeddings_batch) >= self.batch_size:
                self.collection.add(embeddings=embeddings_batch, documents=documents_batch, ids=[str(i) for i in range(i - self.batch_size + 1, i + 1)])
                embeddings_batch = []
                documents_batch = []

        # Ajouter les derniers items restants s'il y en a
        if embeddings_batch:
            self.collection.add(embeddings=embeddings_batch, documents=documents_batch, ids=[str(i) for i in range(len(texts) - len(embeddings_batch), len(texts))])

    def search_context(self, query, n_results=1):
        query_embedding = self.embedding_model.encode([query]).tolist()
        results = self.collection.query(query_embeddings=query_embedding, n_results=n_results)
        return results['documents']


# Initialisation du store de vecteurs et peuplement du dataset
vector_store = VectorStore("embedding_vector")
vector_store.populate_vectors(dataset)

# Fonction pour générer du texte
def generate_text(message, max_tokens, temperature, top_p):
    # Profiler le temps d'exécution de la génération de texte
    start_time = time.time()

    # Récupérer le contexte depuis le store de vecteurs
    context_results = vector_store.search_context(message, n_results=1)
    context = context_results[0] if context_results else ""

    # Créer le modèle de prompt
    prompt_template = (
        f"SYSTEM: You are a recipe generating bot.\n"
        f"SYSTEM: {context}\n"
        f"USER: {message}\n"
        f"ASSISTANT:\n"
    )

    # Générer le texte avec le modèle de langue
    output = llm(
        prompt_template,
        temperature=temperature,
        top_p=top_p,
        top_k=40,
        repeat_penalty=1.1,
        max_tokens=max_tokens,
    )

    # Traiter la sortie
    input_string = output['choices'][0]['text'].strip()
    cleaned_text = input_string.strip("[]'").replace('\\n', '\n')
    continuous_text = '\n'.join(cleaned_text.split('\n'))

    # Afficher le temps d'exécution
    print(f"Temps d'exécution pour générer du texte : {time.time() - start_time} secondes")

    return continuous_text


# Définir l'interface Gradio
demo = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.Textbox(lines=2, placeholder="Enter your message here...", label="Message"),
    ],
    outputs=gr.Textbox(label="Generated Text"),
    title="Chatbot - Your Personal Culinary Advisor",
    description="Running LLM with context retrieval from ChromaDB",
    cache_examples=False,  # Désactivez le cache
    examples=[
        ["I have leftover rice, what can I make out of it?"],
        ["I just have some milk and chocolate, what dessert can I make?"],
        ["I am allergic to coconut milk, what can I use instead in a Thai curry?"],
        ["Can you suggest a vegan breakfast recipe?"],
        ["How do I make a perfect scrambled egg?"],
        ["Can you guide me through making a soufflé?"],
    ],  # Réduire le nombre d'exemples pour accélérer
)

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