import pandas as pd
import streamlit as st
from langchain.docstore.document import Document
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from langchain_community.llms import HuggingFaceHub
import os
from huggingface_hub import HfApi

# Configuration Streamlit
st.set_page_config(page_title="Assistant Support Client", page_icon="🤖", layout="wide")

# Récupération du token depuis les secrets de l'espace
HUGGINGFACE_API_TOKEN = st.secrets["HUGGINGFACE_API_TOKEN"]

# Création du dossier data s'il n'existe pas
DATA_DIR = "data"
os.makedirs(DATA_DIR, exist_ok=True)
EMBEDDINGS_FILE = os.path.join(DATA_DIR, "faq_embeddings.pkl")

# Configuration du client HF Hub
hf_api = HfApi(token=HUGGINGFACE_API_TOKEN)

@st.cache_resource
def load_faq_data():
    """Charge les données FAQ depuis un CSV"""
    try:
        df = pd.read_csv('data/customer_support_faq.csv')
        documents = []
        for _, row in df.iterrows():
            question = row['question']
            answer = row['answer']
            documents.append(Document(page_content=f"Q: {question} A: {answer}"))
        return documents
    except Exception as e:
        st.error(f"Erreur lors du chargement des données FAQ: {str(e)}")
        return []

@st.cache_resource
def initialize_embeddings(_documents):
    """Initialise ou charge les embeddings"""
    try:
        embeddings = HuggingFaceEmbeddings(
            model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
        )
        vector_store = FAISS.from_documents(_documents, embeddings)
        return vector_store
    except Exception as e:
        st.error(f"Erreur lors de l'initialisation des embeddings: {str(e)}")
        return None

@st.cache_resource
def initialize_qa_chain(vector_store):
    """Initialise la chaîne de question-réponse"""
    try:
        # Initialisation du modèle
        llm = HuggingFaceHub(
            repo_id="bigscience/bloom",  # Vous pouvez changer le modèle
            huggingfacehub_api_token=HUGGINGFACE_API_TOKEN,
            model_kwargs={"temperature": 0.5, "max_length": 512}
        )

        # Template de prompt
        prompt_template = """Utilise le contexte suivant pour répondre à la question.
        Si la réponse n'est pas dans le contexte, réponds "Je ne sais pas."
        
        Contexte: {context}
        
        Question: {question}
        
        Réponse:"""

        PROMPT = PromptTemplate(
            template=prompt_template,
            input_variables=["context", "question"]
        )

        # Création de la chaîne QA
        qa_chain = RetrievalQA.from_chain_type(
            llm=llm,
            chain_type="stuff",
            retriever=vector_store.as_retriever(),
            chain_type_kwargs={"prompt": PROMPT},
            return_source_documents=False,
        )
        
        return qa_chain
    except Exception as e:
        st.error(f"Erreur lors de l'initialisation de la chaîne QA: {str(e)}")
        return None

# CSS personnalisé
st.markdown("""
    <style>
    .stApp {
        background-color: #1E1E1E;
    }
    .user_message {
        color: white;
        background-color: #F85046;
        padding: 10px;
        border-radius: 10px;
        margin: 10px 20px 10px auto;
        text-align: left;
        word-wrap: break-word;
        display: inline-block;
        max-width: 80%;
    }
    .bot_message {
        color: #1E1E1E;
        background-color: white;
        padding: 10px;
        border-radius: 10px;
        margin: 10px 20px;
        text-align: left;
        word-wrap: break-word;
        display: inline-block;
        max-width: 80%;
    }
    .header {
        font-size: 40px;
        color: white;
        padding: 20px;
        text-align: center;
        border-radius: 10px;
        font-weight: bold;
    }
    .the_text {
        color: #F9F8FD;
        font-size: 20px;
        margin-top: 20px;
    }
    .the_conv {
        color: #F9F8FD;
        font-size: 30px;
    }
    </style>
""", unsafe_allow_html=True)

# Interface utilisateur
st.markdown('<div class="header">Assistant Support Client</div>', unsafe_allow_html=True)
st.markdown('<div class="the_text">Posez votre question ci-dessous :</div>', unsafe_allow_html=True)

# Initialisation de la session state
if 'conversation' not in st.session_state:
    st.session_state.conversation = []

# Chargement des composants
documents = load_faq_data()
if documents:
    vector_store = initialize_embeddings(documents)
    if vector_store:
        qa_chain = initialize_qa_chain(vector_store)
        if qa_chain:
            # Champ de saisie
            user_query = st.text_input("", key="user_input")

            # Bouton pour obtenir la réponse
            if st.button("Obtenir une réponse", key="submit"):
                if user_query:
                    try:
                        # Obtention de la réponse
                        response = qa_chain({"query": user_query})
                        answer = response["result"]
                        
                        # Ajout à l'historique
                        st.session_state.conversation.append({
                            "user": user_query,
                            "bot": answer
                        })

                        # Affichage de la conversation
                        st.markdown('<div class="the_conv">Conversation :</div>', unsafe_allow_html=True)
                        for chat in st.session_state.conversation:
                            st.markdown(
                                f'<div class="user_message">🤵 : {chat["user"]}</div>',
                                unsafe_allow_html=True
                            )
                            st.markdown(
                                f'<div class="bot_message">🤖 : {chat["bot"]}</div>',
                                unsafe_allow_html=True
                            )
                    except Exception as e:
                        st.error(f"Erreur lors du traitement de la requête: {str(e)}")
                else:
                    st.warning("Veuillez entrer une question.")