import os
import faiss
import pickle
import numpy as np
import re
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, pipeline
from huggingface_hub import InferenceClient

# Token Hugging Face depuis les secrets (Space)
HF_TOKEN = os.environ.get("edup2")
use_client = False

# Tentative de chargement de Mistral
try:
    if HF_TOKEN:
        MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.1"
        client = InferenceClient(MODEL_NAME, token=HF_TOKEN)
        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=HF_TOKEN)
        use_client = True
    else:
        raise ValueError("Pas de token trouvé pour Mistral.")
except Exception as e:
    print(f"⚠️ Impossible de charger Mistral : {e}")
    MODEL_NAME = "google/flan-t5-base"
    generator = pipeline("text2text-generation", model=MODEL_NAME)
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    use_client = False

# Chargement de l’index FAISS et des documents
def load_faiss_index(index_path="faiss_index/faiss_index.faiss", doc_path="faiss_index/documents.pkl"):
    index = faiss.read_index(index_path)
    with open(doc_path, "rb") as f:
        documents = pickle.load(f)
    return index, documents

# Modèle d’embedding
def get_embedding_model():
    return SentenceTransformer("sentence-transformers/multi-qa-MiniLM-L6-cos-v1")

# Recherche dans l’index
def query_index(question, index, documents, model, k=3):
    question_embedding = model.encode([question])
    _, indices = index.search(np.array(question_embedding).astype("float32"), k)
    return [documents[i] for i in indices[0]]

# Nettoyage du contexte
def nettoyer_context(context):
    context = re.sub(r"\[\'(.*?)\'\]", r"\1", context)
    context = context.replace("None", "")
    return context

# Génération de la réponse
def generate_answer(question, context):
    prompt = f"""Voici des informations sur des établissements et formations :

{context}

Formule ta réponse comme un conseiller d’orientation bienveillant, de manière fluide et naturelle.

Question : {question}
Réponse :"""

    if use_client:
        response = client.text_generation(prompt=prompt, max_new_tokens=300, timeout=30)
        return response
    else:
        result = generator(prompt, max_new_tokens=256, do_sample=True)
        return result[0]["generated_text"]