import faiss
import pickle
import numpy as np
import re
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer  # Ajouté pour la gestion des tokens
from huggingface_hub import InferenceClient

# Chargement du modèle Falcon-7B
client = InferenceClient("tiiuae/falcon-7b-instruct")

# Chargement du tokenizer (même base que Falcon)
tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-7b-instruct")

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

def get_embedding_model():
    return SentenceTransformer("sentence-transformers/multi-qa-MiniLM-L6-cos-v1")

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]]

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

def generate_answer(question, context):
    MAX_TOKENS_TOTAL = 2048
    MAX_NEW_TOKENS = 300
    MAX_PROMPT_TOKENS = MAX_TOKENS_TOTAL - MAX_NEW_TOKENS

    # Construction initiale du prompt
    base_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 :"""

    # Tronquer si le prompt est trop long
    tokens = tokenizer.encode(base_prompt)
    if len(tokens) > MAX_PROMPT_TOKENS:
        # Réduction progressive du contexte uniquement
        context_tokens = tokenizer.encode(context)
        keep_tokens = MAX_PROMPT_TOKENS - len(tokenizer.encode(base_prompt.replace(context, "")))
        truncated_context = tokenizer.decode(context_tokens[:keep_tokens])
        
        # Reconstruire le prompt avec contexte réduit
        base_prompt = f"""Voici des informations sur des établissements et formations :

{truncated_context}

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

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

    print("===== PROMPT ENVOYÉ =====")
    print(base_prompt)

    response = client.text_generation(prompt=base_prompt, max_new_tokens=MAX_NEW_TOKENS, timeout=30)

    print("===== RÉPONSE REÇUE =====")
    print(response)

    return response.get("generated_text", response)  # selon format du retour