import streamlit as st
import pandas as pd
from datetime import datetime
from RAG import answer_question, retrieve_chunks, embed_question, build_context

# Configuration de la page
st.set_page_config(
    page_title="Tindle RAG Chat", 
    page_icon="💬", 
    layout="wide",
    initial_sidebar_state="expanded"
)

# CSS personnalisé pour le thème
st.markdown("""
<style>
    .stApp {
        background-color: var(--background-color);
    }
    .chat-message {
        padding: 1rem;
        border-radius: 0.5rem;
        margin-bottom: 1rem;
        border: 1px solid var(--border-color);
    }
    .source-preview {
        background-color: var(--secondary-background-color);
        padding: 0.5rem;
        border-radius: 0.25rem;
        margin: 0.5rem 0;
        border-left: 3px solid #1f77b4;
    }
    .export-button {
        margin: 1rem 0;
    }
</style>
""", unsafe_allow_html=True)

# Sidebar pour les paramètres
with st.sidebar:
    st.title("⚙️ Paramètres")
    
    # Sélecteur de thème
    theme = st.selectbox(
        "🎨 Thème",
        ["Clair", "Sombre"],
        index=0
    )
    
    # Appliquer le thème
    if theme == "Sombre":
        st.markdown("""
        <style>
            .stApp {
                background-color: #0e1117;
                color: #fafafa;
            }
            .stChatMessage {
                background-color: #262730;
            }
        </style>
        """, unsafe_allow_html=True)
    
    # Paramètres RAG
    st.subheader("🔍 Paramètres de recherche")
    top_k = st.slider("Nombre de passages à récupérer", 5, 20, 10)
    show_sources = st.checkbox("Afficher les sources", value=True)
    
    # Statistiques
    st.subheader("📊 Statistiques")
    if "messages" in st.session_state:
        st.metric("Questions posées", len([m for m in st.session_state["messages"] if m["role"] == "user"]))
    
    # Bouton de réinitialisation
    if st.button("🗑️ Effacer l'historique"):
        st.session_state["messages"] = []
        st.session_state["sources"] = []
        st.rerun()

# Titre principal
st.title("💬 Tindle RAG Chat")
st.markdown("Posez une question sur le droit fiscal, l'assistant RAG vous répondra en s'appuyant sur les sources indexées.")

# Initialisation des variables de session
if "messages" not in st.session_state:
    st.session_state["messages"] = []
if "sources" not in st.session_state:
    st.session_state["sources"] = []

# Affichage de l'historique du chat avec sources
for i, msg in enumerate(st.session_state["messages"]):
    if msg["role"] == "user":
        with st.chat_message("user"):
            st.markdown(msg["content"])
    else:
        with st.chat_message("assistant"):
            st.markdown(msg["content"])
            
            # Afficher les sources si disponibles et demandées
            # Correction de l'index : on cherche les sources pour cette réponse spécifique
            source_index = len([m for m in st.session_state["messages"][:i] if m["role"] == "assistant"])
            if show_sources and source_index < len(st.session_state["sources"]) and st.session_state["sources"][source_index]:
                with st.expander(f"📚 Sources utilisées ({len(st.session_state['sources'][source_index])} passages)"):
                    for j, source in enumerate(st.session_state["sources"][source_index], 1):
                        # Construire la section métadonnées
                        metadata_parts = []
                        for key, value in source["metadata"].items():
                            metadata_parts.append(f"{key}: {value}")
                        
                        metadata_str = f" | ".join(metadata_parts) if metadata_parts else ""
                        source_info = f"<strong>Source {j}:</strong> {source['id']}"
                        if metadata_str:
                            source_info += f" | {metadata_str}"
                        source_info += f" <em>(score: {source['score']:.3f})</em>"
                        
                        st.markdown(f"""
                        <div class="source-preview">
                        {source_info}<br>
                        {source['text']}
                        </div>
                        """, unsafe_allow_html=True)

# Saisie utilisateur avec indicateurs de chargement détaillés
if prompt := st.chat_input("Votre question..."):
    st.session_state["messages"].append({"role": "user", "content": prompt})
    
    # Affichage des étapes de traitement
    progress_bar = st.progress(0)
    status_text = st.empty()
    
    try:
        # Étape 1: Embedding
        status_text.text("🔄 Création de l'embedding de la question...")
        progress_bar.progress(25)
        q_emb = embed_question(prompt)
        
        # Étape 2: Récupération des passages
        status_text.text("🔍 Recherche des passages pertinents...")
        progress_bar.progress(50)
        top_chunks = retrieve_chunks(q_emb, top_k)
        st.session_state["sources"].append(top_chunks)
        
        # Étape 3: Construction du contexte
        status_text.text("📝 Construction du contexte...")
        progress_bar.progress(75)
        context = build_context(top_chunks)
        
        # Étape 4: Génération de la réponse
        status_text.text("🤖 Génération de la réponse...")
        progress_bar.progress(90)
        response = answer_question(prompt, top_k)
        
        # Finalisation
        progress_bar.progress(100)
        status_text.text("✅ Réponse générée avec succès!")
        
        st.session_state["messages"].append({"role": "assistant", "content": response})
        
    except Exception as e:
        status_text.text("❌ Erreur lors du traitement")
        response = f"Erreur lors de l'appel au RAG : {e}"
        st.session_state["messages"].append({"role": "assistant", "content": response})
        st.session_state["sources"].append([])
    
    # Nettoyage des indicateurs
    progress_bar.empty()
    status_text.empty()
    
    st.rerun()

# Footer avec informations
st.markdown("---")
st.markdown("""
<div style='text-align: center; color: #666;'>
    <small>
        💡 Conseil: Plus votre question est précise, plus la réponse sera pertinente.<br>
        🔍 Les sources utilisées sont affichées pour chaque réponse.
    </small>
</div>
""", unsafe_allow_html=True)