# import os
# from dotenv import load_dotenv
# from langchain_community.chat_models import ChatOpenAI
import streamlit as st
from chats.streamlit_tools import import_file  # ,clear_cache
from streamlit_extras.add_vertical_space import add_vertical_space
from langchain_tools.pdf_tools import PdfLangChain
from langchain_tools.lc_tools import LangChainTools
from chats.chat_tools import MessageManager


# App title
st.set_page_config(page_title="LLMOneClusterTeam")

# sidebar
with st.sidebar:
    # Cargar el logo (asegúrate de que el archivo de imagen esté en la misma carpeta que tu script)
    logo_path = "documents/Logo azulblanco.png"
    logo = st.sidebar.image(logo_path, width=200)

    # Ajusta el ancho según sea necesario
    add_vertical_space(28)
    # pdf_name = import_file()
    st.markdown("Built by [OneCluster](https://www.onecluster.org/).")


col1, col2 = st.columns([1.1, 1])
with col1:
    st.title(
        "DocumentAssist",
    )
with col2:
    logo_2 = st.image("documents/pdfs/logo_1-removebg-preview.png", width=110)

pdf_name = import_file()
if pdf_name:
    with st.spinner("Processing the document..."):
        # Inicializamos la clase PdfLangChain
        pdfLangChain = PdfLangChain(pdf_name)
        pdf_name = pdfLangChain.file_name

        # Cargamos el documento PDF
        docs: list = pdfLangChain.load_pdf()

        # Dividimos los documentos en partes mas pequenas
        docs_split: list = pdfLangChain.split_docs(docs)

        # Instanciamos la clase LangChainTools que contiene herramientras LangChain
        langChainTools = LangChainTools()

        # Cargamos el modelo de embeddings
        embedding_model = langChainTools.load_embedding_opnai()

        # Creamos el vector store
        docstorage = langChainTools.create_vector_strore(
            docs_split, pdf_name, embedding_model
        )

    # Cargamos el modelo LLM desde LangChain
    llm = langChainTools.load_llm_open_source()

    # Creamos la cadena que integra Vectorstroe, el LLM para hacer consultas.
    # Para este caso la cadena tene el parametro de memoria.
    qa = langChainTools.define_retrieval_qa_memory(
        llm, docstorage, pdf_name, embedding_model
    )

    # Store conversation history
    if "messages" not in st.session_state.keys():
        st.session_state.messages = [
            {
                "role": "assistant",
                "content": "Hola, soy una IA con el que puedes chatear con tu PDF. Haz un pregunta al documento.",
            }
        ]

    # Display or clear chat messages
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.write(message["content"])

    def clear_chat_history():
        st.session_state.messages = [
            {
                "role": "assistant",
                "content": "Hola, soy una IA con el que puedes chatear con tu PDF. Haz un pregunta al documento.",
            }
        ]

    st.sidebar.button("Clear chat history", on_click=clear_chat_history)

    @st.cache_resource
    def get_num_tokens(prompt):
        """Get the number of tokens in a given prompt"""
        return len(prompt.split())

    # Function for generating Snowflake Arctic response

    # User-provided prompt
    if prompt := st.chat_input():
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.write(prompt)

        # Generate a new response if last message is not from assistant
        if st.session_state.messages[-1]["role"] != "assistant":
            with st.spinner("Thinking..."):
                # Creamos la cadena que integra Vectorstroe, el LLM para hacer consultas.
                # Para este caso la cadena tene el parametro de memoria.
                qa = langChainTools.define_retrieval_qa_memory(
                    llm, docstorage, pdf_name, embedding_model
                )

                input = "\n".join([msg["content"] for msg in st.session_state.messages])

                query = qa.invoke({"question": f"{prompt}"}, return_only_outputs=True)

                response_text = query["answer"]
                documents_source = query["source_documents"]

                messageManager = MessageManager()

                citation: str = messageManager.generate_citations(documents_source)
                # st.markdown(citation)

            with st.chat_message("assistant"):
                st.write(response_text)
                st.session_state.messages.append(
                    {"role": "assistant", "content": response_text}
                )
                expander = st.expander("Fuentes")
                expander.markdown(citation)