import streamlit as st
from graph_agentA import agent as agent_A
from graph_agentB import agent as agent_B
from graph_agentC import agent as agent_C
from config import *
from dotenv import load_dotenv
from pinecone_utilsB import *
from langchain_core.messages import AIMessageChunk
from typing import Literal
import re

from carbon.empreinte_carbone import initCarbon, traceImpact, display_cf_comparison


# Charger les variables d'environnement
load_dotenv()

# Initialiser l'état de chat globalement
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

def check_indexes_ready():
    """Vérifie que les index Pinecone sont prêts."""
    try:
        # Vérifier que les index existent
        existing_indexes = pc.list_indexes().names()
        if sparse_index_name not in existing_indexes or dense_index_name not in existing_indexes:
            st.error("Les index Pinecone ne sont pas prêts. Veuillez exécuter 'initialize_indexes.py'.")
            return False
        return True
    except Exception as e:
        st.error(f"Erreur lors de la vérification des index Pinecone : {e}")
        return False

def process_query(query, architecture: Literal["A", "B", "C"]):
    """Traite la requête de l'utilisateur avec l'architecture A, B ou C."""

    # Reload conversation
    display_chat_history()

    config = {"metadata": {"architecture": architecture}, "tags": ["arch_" + architecture]}
    # Récupération des paramètres dynamiques uniquement
    k = st.session_state.get("k", 30)  # Nombre de documents
    similarity_threshold = st.session_state.get("similarity_threshold", 0.7)  # similarité cosinus

    if architecture == "A":
        agent = agent_A
        initial_state = {
            "query": query,
            "messages": [],
            "relevant_docs": [],
            "response": "",
            "k": k,
            "similarity_threshold": similarity_threshold,
        }

    elif architecture in ["B", "C"]:
        agent = agent_B if architecture == "B" else agent_C

        # Récupération du paramétre alpha uniquement pour B et C
        alpha = st.session_state.get("alpha", 0.5)  # Pondération hybride
        
        initial_state = {
            "query": query,
            "messages": [],
            "relevant_docs": [],
            "response": "",
            "k": k,  
            "alpha": alpha,  
            "similarity_threshold": similarity_threshold,
        }

    st.session_state.chat_history.append({"role": "user", "content": query})

    with st.chat_message("user"):
        st.markdown(query)

    full_response = ""
    usages = None
    start_time = time.time()  # Start timing

    events = agent.stream(initial_state, config=config, stream_mode="messages")

    # Ajouter le message du chatbot avec streaming
    with st.chat_message("assistant"):
        response_placeholder = st.empty()

        for event in events:
            for message in event:
                if isinstance(message, AIMessageChunk):
                    if hasattr(message, 'content'):
                        full_response += message.content

                        # Supprimer les requêtes Cypher avant l'affichage
                        full_response = re.sub(r"(?i)(MATCH|CREATE|MERGE|DELETE|CALL)[\s\S]+?;", "", full_response).strip()

                        response_placeholder.markdown(full_response)

                    if hasattr(message, 'usage_metadata'):
                        usages = message.usage_metadata

        latency = time.time() - start_time  # Calculate elapsed time

        #         if isinstance(message, AIMessageChunk) and hasattr(message, 'content'):
        #             full_response += message.content  # Accumuler les morceaux
        
        
        # 🛑 Vérification si la réponse contient une requête Cypher
        if re.search(r"(?i)(MATCH|CREATE|MERGE|DELETE|CALL)[\s\S]+;", full_response):
            
            full_response = full_response = re.sub(r"(?i)(MATCH|CREATE|MERGE|DELETE|CALL)[\s\S]+?;", "", full_response).strip()


        # 🔄 Affichage final
        # response_placeholder.markdown(full_response)
        st.session_state.chat_history.append({"role": "assistant", "content": full_response})

        traceImpact(
            st.session_state.current_model,
            trace={
                'completion_tokens' : usages.get('output_tokens', None),
                'latency' : latency  # Use the measured latency
            }
        )
        
        # Mise à jour des tokens + coût
        st.session_state['tokens_metrics']['input_tokens'] += usages.get('input_tokens', 0)
        st.session_state['tokens_metrics']['output_tokens'] += usages.get('output_tokens', 0)
        st.session_state['tokens_metrics']['total_tokens'] += usages.get('total_tokens', 0)
        calculate_tokens_cost()

        st.rerun()
        
    

def display_sidebar():
    """Affiche la barre latérale."""
    with st.sidebar:
        #st.title("📄 La confession muette")
        #st.write("Posez vos questions sur le document.")
        
        st.markdown("### Document de référence 📄")
        st.markdown("<h2 style='text-align: center;margin-top:20px;margin-bottom:20px;'>La confession muette (2025)</h2>", unsafe_allow_html=True)


        lien_ressource = "https://www.fnac.com/livre-numerique/a21290809/Gaspard-Boreal-La-Confession-muette"
        st.markdown("*Avec l'aimable autorisation de Gaspard Boréal :* \n[*Récit d'origine*]({})".format(lien_ressource))


        # Token metrics containers
        # st.sidebar.markdown("### Tokens")

        st.markdown("### Paramètres de la recherche RAG")

        # Sélection du nombre de documents (k)
        st.markdown("""
        **Nombre de documents à récupérer (k)**  
        <small>🛈 Détermine combien de documents seront récupérés lors de la recherche.</small>
        """, unsafe_allow_html=True)

        st.number_input(" ", min_value=1, max_value=100, value=30, step=1, key="k")

        # Sélection du score de similarité cosinus
        st.markdown("""
        **Score de similarité cosinus**  
        <small>🛈 Ce paramètre définit le seuil minimal de similarité entre deux vecteurs. Plus il est élevé, plus seuls les éléments très similaires seront considérés comme correspondants.</small>
        """, unsafe_allow_html=True)
        st.slider(" ", 0.0, 1.0, value=0.7, step=0.05, key="similarity_threshold")

        # Afficher alpha uniquement pour B et C
        if st.session_state.get("architecture") in ["B", "C"]:
            st.markdown("""
            **Équilibre entre recherche sémantique et syntaxique**  
            <small>🛈 Pour `alpha = 0.0`, la recherche est purement syntaxique. Pour `alpha = 1.0`, elle est purement sémantique.</small>
            """, unsafe_allow_html=True)

            st.slider(" ", 0.0, 1.0, value=0.5, step=0.05, key="alpha")  # Sauvegarde de la valeur alpha
        else:
            # Réinitialiser alpha si l'architecture est A
            st.session_state['alpha'] = None  


        st.sidebar.markdown("### API Mistral AI")

        # HTML pour le tableau sans index ni colonnes inutiles
        table_html = f"""
            <table style="width:100%">
                <tr><td>Input Tokens</td><td style="text-align: right;"><b>{st.session_state['tokens_metrics'].get('input_tokens', 0)}</b></td></tr>
                <tr><td>Output Tokens</td><td style="text-align: right;"><b>{st.session_state['tokens_metrics'].get('output_tokens', 0)}</b></td></tr>
                <tr><td>Total Tokens</td><td style="text-align: right;"><b>{st.session_state['tokens_metrics'].get('total_tokens', 0)}</b></td></tr>
                <tr><td>Coût</td><td style="text-align: right;"><b>{st.session_state['tokens_cost']} €</b></td></tr>
            </table>
        """

        st.sidebar.markdown(table_html, unsafe_allow_html=True)
        
        st.sidebar.markdown("### Empreinte Carbone")
        
        display_cf_comparison(st.sidebar)
        
        
        st.sidebar.markdown("---")

        st.sidebar.markdown("### Documentation")
        st.markdown("*Livre de référence :* La confession muette (2025)")
        st.markdown("*Test Q&A :* [*Lien*]({})".format("https://docs.google.com/spreadsheets/d/1dkMgv1MM9ZQjY8EIsR3V8DpMjlmSG37ydFD9TyTHtt4/edit?usp=sharing"))
        st.markdown("*Documentation RAG :* [*Lien*]({})".format("https://docs.google.com/presentation/d/17Q3qt0-p9feSxSY1AgLjoDcitwg_TfNlnzDSynkt8ew/edit?usp=sharing"))

        
       
        st.sidebar.markdown("---")
        st.sidebar.image("./assets/bziiit.png", width=100)

        st.markdown("• *Contributeur AFNOR SPEC IA FRUGAL*")
        st.markdown("• *Certifié AFNOR COMPETENCES*") 
        st.markdown("• *Membre HUB FRANCE IA*")
        st.markdown("• *Membre OSFARM*")

        st.sidebar.markdown("")

        st.sidebar.markdown("2025 : Open source en Licence MIT")
        st.sidebar.markdown("info@bziiit.com")

def display_chat_history():
    """Affiche l'historique de chat."""
    for message in st.session_state.chat_history:
        if message["role"] == "user":
            with st.chat_message("user"):
                st.markdown(message["content"])
        elif message["role"] == "assistant":
            with st.chat_message("assistant"):
                st.markdown(message["content"])

def is_neo4j_aura_active(retries: int, wait_seconds: int):
    """
    Tente de se connecter à l'instance Neo4j. Si elle est en veille,
    la fonction essaiera de se reconnecter jusqu'à ce que l'instance soit réveillée.

    Args:
        retries (int): Nombre de tentatives avant d'abandonner.
        wait_seconds (int): Temps d'attente entre les tentatives.
    """
    for attempt in range(1, retries + 1):
        try:
            with neo4j_driver.session() as session:
                session.run("RETURN 1") # Simple requête pour réveiller l’instance
            return True
        except Exception as e:
            print(f"⚠️ Tentative {attempt}/{retries} : l'instance est en veille. Nouvel essai dans {wait_seconds}s...")
            time.sleep(wait_seconds)
    raise RuntimeError("❌ Impossible de réveiller l'instance Neo4j après plusieurs tentatives.")

def main():
    initCarbon()
    st.session_state['current_model'] = "mistral-large-latest"

    if "tokens_metrics" not in st.session_state or "chat_history" not in st.session_state:
        initialize_conversation()

    if not check_indexes_ready():
        return

    st.title("LaConfessionMuette Chatbot")
    st.markdown("""
        <style>
        div[data-testid="stSelectbox"] {
            width: 200px !important;  
        }
        </style>
        """, unsafe_allow_html=True)

    st.markdown("• Basique : *Recherche sémantique*")
    st.markdown("• Intermédiaire : *Rechercher hybride (sémantique + mots clés)*") 
    st.markdown("• Avancée : *Graph de connaissance + Recherche hybride*")

    # Initialisation des états
    if "architecture" not in st.session_state:
        st.session_state["architecture"] = "A"
    if "selected_arch" not in st.session_state:
        st.session_state["selected_arch"] = "Basic"
    if "first_time_advanced" not in st.session_state:
        st.session_state["first_time_advanced"] = True
    if "advanced_info_shown" not in st.session_state:
        st.session_state["advanced_info_shown"] = False

    # Selectbox
    new_selection = st.selectbox(
        "Sélectionnez une architecture RAG :",
        ["Basic", "Intermédiaire", "Avancée"],
        index=["Basic", "Intermédiaire", "Avancée"].index(st.session_state["selected_arch"]),
        key="arch_selection"
    )

    # Si la sélection change
    if new_selection != st.session_state["selected_arch"]:
        if new_selection == "Avancée":
            # Message de confirmation pour première fois
            if st.session_state["first_time_advanced"]:
                st.info("⚠️ Attention, ce type de fonctionnement est plus précis mais consomme plus de carbone. Souhaitez-vous réellement l'utiliser ?")
                col1, col2 = st.columns(2)
                if col1.button("✅ Confirmer"):
                    st.session_state["architecture"] = "C"
                    st.session_state["selected_arch"] = "Avancée"
                    st.session_state["first_time_advanced"] = False
                    st.session_state["advanced_info_shown"] = True
                    st.toast("Mode avancé activé")
                    st.rerun()
                if col2.button("❌ Annuler"):
                    st.session_state["selected_arch"] = "Intermédiaire"
                    st.rerun()
                return
            else:
                st.session_state["architecture"] = "C"
                st.session_state["selected_arch"] = "Avancée"
                st.session_state["advanced_info_shown"] = True
        else:
            st.session_state["architecture"] = "A" if new_selection == "Basic" else "B"
            st.session_state["selected_arch"] = new_selection
        st.rerun()

    # Affichage du message d'information persistant
    if st.session_state["selected_arch"] == "Avancée" and st.session_state["advanced_info_shown"]:
        st.info("ℹ️ Vous utilisez actuellement l'architecture avancée. Cette option est optimisée pour la précision mais nécessite davantage de ressources.")
  

    display_sidebar()

    if st.session_state.chat_history:
        display_chat_history()
    
    query = st.chat_input("Posez votre question ici:")
    if query:
        if st.session_state.get("architecture") == "C":
            #with st.spinner("⏳ Activation de l'instance Neo4j AuraDB en cours..."):
                if is_neo4j_aura_active(3, 10):  # 3 retries, 10 seconds each
                    st.toast("✅ Neo4j AuraDB est actif.")
                    process_query(query, "C")
                else:
                    st.error("❌ Échec de l'activation de Neo4j. Veuillez réessayer plus tard.")
        else:
            process_query(query, st.session_state["architecture"])

           
def calculate_tokens_cost():
    input_tokens = st.session_state['tokens_metrics'].get('input_tokens', 0)
    output_tokens = st.session_state['tokens_metrics'].get('output_tokens', 0)
    cost_input = (input_tokens / 1_000_000) * 1.80
    cost_output = (output_tokens / 1_000_000) * 5.40
    total_cost = round(cost_input + cost_output, 3)
    st.session_state['tokens_cost'] = total_cost
            
def initialize_conversation():
    # st.session_state["messages"] = []
    st.session_state["chat_history"] = []
    st.session_state["tokens_cost"] = 0
    st.session_state['tokens_metrics'] = {
        'input_tokens': 0,
        'output_tokens': 0,
        'total_tokens': 0
    }
    if "k" not in st.session_state:
        st.session_state['k'] = 30  # Valeur par défaut pour k
    if "similarity_threshold" not in st.session_state:
        st.session_state['similarity_threshold'] = 0.7  # Valeur par défaut pour similarity_threshold
    if "alpha" not in st.session_state:
        st.session_state['alpha'] = 0.5 
    calculate_tokens_cost()



if __name__ == "__main__":
    main()