import streamlit as st
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langchain.prompts import PromptTemplate
from model import selector
from util import getYamlConfig
from st_copy_to_clipboard import st_copy_to_clipboard

def display_messages():

    for i, message in enumerate(st.session_state.chat_history):
        if isinstance(message, AIMessage):
            with st.chat_message("AI"):
                # Display the model from the kwargs
                model = message.kwargs.get("model", "Unknown Model")  # Get the model, default to "Unknown Model"
                st.write(f"**Model :** {model}")
                st.markdown(message.content)
                st_copy_to_clipboard(message.content,key=f"message_{i}")
                # show_retrieved_documents(st.session_state.chat_history[i-1].content)
        
        elif isinstance(message, HumanMessage):
            with st.chat_message("Moi"):
                st.write(message.content)

        elif isinstance(message, SystemMessage):
            with st.chat_message("System"):
                st.write(message.content)

def show_retrieved_documents(query: str = ''):
    if query == '':
        return
    
    # Créer l'expander pour les documents trouvés
    expander = st.expander("Documents trouvés")
    
    # Boucler à travers les documents récupérés
    for item in st.session_state.get("retrived_documents", []):
        if 'query' in item:
            if item["query"] == query:
                for doc in item.get("documents", []):
                    expander.write(doc["metadata"]["source"])


def launchQuery(query: str = None):

    # Initialize the assistant's response
    full_response = st.write_stream(
        st.session_state["assistant"].ask(
            query,
            # prompt_system=st.session_state.prompt_system,
            messages=st.session_state["chat_history"] if "chat_history" in st.session_state else [],
            variables=st.session_state["data_dict"]
        ))

    # Temporary placeholder AI message in chat history
    st.session_state["chat_history"].append(AIMessage(content=full_response, kwargs={"model": st.session_state["assistant"].getReadableModel()}))
    st.rerun()


def show_prompts():
    yaml_data = getYamlConfig()["prompts"]
    
    expander = st.expander("Prompts pré-définis")
    
    for categroy in yaml_data:
        expander.write(categroy.capitalize())

        for item in yaml_data[categroy]:
            if expander.button(item, key=f"button_{item}"):
                launchQuery(item)

def remplir_texte(texte: str, variables: dict, remove_line_if_unspecified: bool = False) -> str:
    # Convertir les valeurs en chaînes de caractères pour éviter les erreurs avec format()
    variables_str = {
        key: (', '.join(value) if isinstance(value, list) and len(value) else value if value else 'Non spécifié')
        for key, value in variables.items()
    }

    # Remplacer les variables dynamiques dans le texte
    try:
        texte_rempli = texte.format(**variables_str)
    except KeyError as e:
        raise ValueError(f"Clé manquante dans le dictionnaire : {e}")
    
    # Supprimer les lignes contenant "Non spécifié" si l'option est activée
    if remove_line_if_unspecified:
        lignes = texte_rempli.split('\n')
        lignes = [ligne for ligne in lignes if 'Non spécifié' not in ligne]
        texte_rempli = '\n'.join(lignes)
    
    return texte_rempli

def page():
    st.subheader("Posez vos questions")

    if "assistant" not in st.session_state:
        st.text("Assistant non initialisé")

    if "chat_history" not in st.session_state or len(st.session_state["chat_history"]) < 2:

        if st.session_state["data_dict"] is not None:
            # Convertir la liste en dictionnaire avec 'key' comme clé et 'value' comme valeur
            vars = {item['key']: item['value'] for item in st.session_state["data_dict"] if 'key' in item and 'value' in item}

        system_template = st.session_state.prompt_system
        full = remplir_texte(system_template, vars, st.session_state["remove_undefined_value"])

        st.session_state["chat_history"] = [
            SystemMessage(content=full),
        ]

    st.markdown("<style>iframe{height:50px;}</style>", unsafe_allow_html=True)

    # Collpase for default prompts
    show_prompts()

    # Models selector
    selector.ModelSelector()

    if(len(st.session_state["chat_history"])):
        if st.button("Effacer l'historique"):
            st.session_state["chat_history"] = []

    # Displaying messages
    display_messages()


    user_query = st.chat_input("")
    if user_query is not None and user_query != "":

        st.session_state["chat_history"].append(HumanMessage(content=user_query))
        
        # Stream and display response
        launchQuery(user_query)


page()