import streamlit as st
from langchain_community.document_loaders import TextLoader, PyMuPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_google_genai import GoogleGenerativeAIEmbeddings, ChatGoogleGenerativeAI
from langchain.schema import SystemMessage, HumanMessage
import requests
import time

# Clé API Google Gemini
API_KEY = "AIzaSyAVUg8B8r-oddxEAzonh3g_7KOyeL3wUK0"

# Charger et traiter les fichiers
pdf_loader = PyMuPDFLoader("doc_0.pdf")
txt_loader = TextLoader("UCAD info3.txt")

pdf_docs = pdf_loader.load()
txt_docs = txt_loader.load()
documents = pdf_docs + txt_docs  # Fusionner les documents

# Découpage des textes en chunks
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
chunks = text_splitter.split_documents(documents)

# Génération des embeddings et indexation
embedding_model = GoogleGenerativeAIEmbeddings(model="models/embedding-001", google_api_key=API_KEY)
vector_store = FAISS.from_documents(chunks, embedding_model)

# Modèle de génération de texte
chat_model = ChatGoogleGenerativeAI(model="gemini-1.5-flash", api_key=API_KEY)

# 🔹 Configuration de la page en mode large
st.set_page_config(page_title="Assistant UCAD (Université Cheikh Anta Diop de Dakar)", page_icon="💡", layout="wide")

# 🔹 Initialisation des messages
if "messages" not in st.session_state:
    st.session_state.messages = [
        {"role": "bot", "content": "Bonjour, je suis l'assistant UCAD ! Comment puis-je vous aider ? 😊"}
    ]

# 🔹 Affichage des messages dans un conteneur
st.header("Assistant UCAD (Université Cheikh Anta Diop de Dakar)")

for msg in st.session_state.messages:
    role_prefix = "🤖 Assistant UCAD" if msg["role"] == "bot" else "👤 Vous"
    st.markdown(f"**{role_prefix}** : {msg['content']}")

# 🔹 Zone de saisie utilisateur
user_input = st.text_input("Posez votre question ici...")

if st.button("Envoyer"):
    if user_input:
        st.session_state.messages.append({"role": "user", "content": user_input})

        # 🔹 Affichage du spinner (chargement)
        with st.spinner("L'Assistant réfléchit... 🤔"):
            time.sleep(1.5)  # Simule un délai de réponse

        # 🔹 Recherche des documents pertinents avec FAISS
        retrieved_docs = vector_store.similarity_search(user_input, k=3)
        context = "\n\n".join([doc.page_content for doc in retrieved_docs])

        # 🔹 Envoi de la question à Google Gemini pour la génération de réponse
        messages = [
            SystemMessage(content="Tu es un assistant qui répond aux questions en utilisant les informations fournies."),
            HumanMessage(content=f"Document : {context}\n\nQuestion : {user_input}\nRéponse détaillée :")
        ]

        try:
            response = chat_model(messages)
            bot_response = response.content
        except Exception:
            bot_response = "Je n'ai pas pu obtenir une réponse pour le moment."

        # 🔹 Ajouter la réponse du bot
        st.session_state.messages.append({"role": "bot", "content": bot_response})

        # 🔄 Rafraîchir l'affichage
        st.rerun()