import streamlit as st
import os
from llama_cpp import Llama
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from huggingface_hub import hf_hub_download

# --- Configuration de la page Streamlit ---
st.set_page_config(page_title="Votre Coach RAG", layout="wide")
st.title("Votre Coach Expert")
st.write("Posez une question sur vos documents, et je vous répondrai en me basant sur leur contenu.")

# --- Fonctions de chargement mises en cache ---
# @st.cache_resource est CRUCIAL pour que Streamlit ne recharge pas les modèles à chaque interaction

@st.cache_resource
def load_llm():
    # ✅ MODIFICATION : On télécharge le modèle depuis le Hub au lieu de le chercher localement.
    # Cela contourne la limite de stockage de 1 Go du Space.
    model_repo_id = "QuantFactory/Meta-Llama-3-8B-Instruct-GGUF"
    model_filename = "Meta-Llama-3-8B-Instruct.Q4_K_M.gguf"
    
    with st.spinner(f"Téléchargement du modèle '{model_filename}'... (Cette étape est longue et n'a lieu qu'une seule fois)"):
        # Télécharge le fichier s'il n'est pas dans le cache et retourne son chemin
        model_path = hf_hub_download(repo_id=model_repo_id, filename=model_filename)
    
    with st.spinner("Chargement du modèle LLM en mémoire..."):
        # ✅ MODIFICATION : n_gpu_layers=0 car nous utilisons le CPU gratuit.
        llm = Llama(model_path=model_path, n_gpu_layers=0, n_ctx=4096, verbose=False, chat_format="llama-3")
    return llm

@st.cache_resource
def load_retriever(faiss_path, embeddings_path):
    with st.spinner("Chargement de la base de connaissances (FAISS)..."):
        # ✅ MODIFICATION : On spécifie 'cpu' car nous n'avons pas de GPU.
        embeddings_model = HuggingFaceEmbeddings(model_name=embeddings_path, model_kwargs={'device': 'cpu'})
        vectorstore = FAISS.load_local(faiss_path, embeddings_model, allow_dangerous_deserialization=True)
    return vectorstore.as_retriever(search_kwargs={"k": 5})

# --- Chemins d'accès (relatifs) ---
DOSSIER_PROJET = os.path.dirname(__file__)
CHEMIN_INDEX_FAISS = os.path.join(DOSSIER_PROJET, "faiss_index_wize")
CHEMIN_MODELE_EMBEDDINGS = os.path.join(DOSSIER_PROJET, "embedding_model")

# --- Chargement des modèles via Streamlit ---
try:
    llm = load_llm()
    retriever = load_retriever(CHEMIN_INDEX_FAISS, CHEMIN_MODELE_EMBEDDINGS)
    st.success("Les modèles sont chargés et prêts !")
except Exception as e:
    st.error(f"Erreur lors du chargement des modèles : {e}")
    st.stop()

# --- Initialisation de l'historique de chat ---
if "messages" not in st.session_state:
    st.session_state.messages = []

# Afficher les messages de l'historique
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# --- Logique de Chat ---
if prompt := st.chat_input("Posez votre question ici..."):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        with st.spinner("Je réfléchis..."):
            # 1. Récupérer le contexte
            docs = retriever.invoke(prompt)
            context = "\n".join([doc.page_content for doc in docs])

            # 2. Créer le prompt pour le LLM
            system_prompt = "Vous êtes Un coach expert. Répondez à la question en vous basant uniquement sur le contexte fourni."
            full_prompt = f"""<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n{system_prompt}\nContexte : {context}<|eot_id|><|start_header_id|>user<|end_header_id|>\nQuestion : {prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>"""
            
            # 3. Générer la réponse
            response = llm(full_prompt, max_tokens=1500, stop=["<|eot_id|>"], echo=False)
            answer = response['choices'][0]['text'].strip()
            st.markdown(answer)
    
    st.session_state.messages.append({"role": "assistant", "content": answer})