import streamlit as st
import os
from langchain.document_loaders import PyPDFDirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.embeddings import SentenceTransformerEmbeddings
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain import hub
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

# Configuración inicial y variables de entorno
st.title("Busca en pdfs")

# Entradas de usuario para las API keys
SMITH_APIKEY = st.text_input('Smith Api Key', type='password')
GOOGLE_API_KEY = st.text_input('Gemini Api Key', type='password')
gemini_model = st.selectbox('Selecciona el modelo de Gemini', ['gemini-1.5-pro-001', 'gemini-1.5-flash-001'])

if SMITH_APIKEY and GOOGLE_API_KEY:
    os.environ["LANGCHAIN_TRACING_V2"] = "true"
    os.environ["LANGCHAIN_API_KEY"] = SMITH_APIKEY

    # Función principal
    def main():
        st.write("Aplicación en Streamlit ejecutándose correctamente")

        # Configura las carpetas necesarias
        source_data_folder = "MisDatos"
        path_db = "VectorDB"
        os.makedirs(source_data_folder, exist_ok=True)
        os.makedirs(path_db, exist_ok=True)

        # Carga de archivos PDF
        uploaded_files = st.file_uploader("Sube archivos PDF", type="pdf", accept_multiple_files=True)
        if uploaded_files:
            for uploaded_file in uploaded_files:
                with open(os.path.join(source_data_folder, uploaded_file.name), "wb") as f:
                    f.write(uploaded_file.getbuffer())
            st.write(f"{len(uploaded_files)} archivos subidos exitosamente")

            # Leer y procesar PDFs
            loader = PyPDFDirectoryLoader(source_data_folder)
            data_on_pdf = loader.load()
            st.write(f"Se han cargado {len(data_on_pdf)} documentos")

            # Particionar datos
            text_splitter = RecursiveCharacterTextSplitter(
                separators=["\n\n", "\n", ". ", " ", ""],
                chunk_size=1024,
                #chunk_size=1000,
                chunk_overlap=200
            )
            splits = text_splitter.split_documents(data_on_pdf)
            st.write(f"Se han creado {len(splits)} fragmentos")

            # Embeddings
            #embeddings_model = SentenceTransformerEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
            embeddings_model = SentenceTransformerEmbeddings(model_name="BAAI/bge-base-en-v1.5")

            # Base de datos vectorial
            #vectorstore = Chroma.from_documents(documents=splits, embedding=embeddings_model.embed_documents, persist_directory=path_db)
            vectorstore = Chroma.from_documents(documents=splits, embedding=embeddings_model, persist_directory=path_db)

            # Configuración del LLM
            system_prompt = """
                            Eres un experto en análisis de datos. Tu tarea es analizar datos, interpretar resultados y proporcionar información valiosa a partir de conjuntos de datos.
                            Tienes un profundo conocimiento de estadística,
                            visualización de datos, minería de datos y técnicas de aprendizaje automático. Siempre debes proporcionar respuestas claras,
                            concisas y basadas en hechos.
                            Das respuestas en español.
                            """
            llm = ChatGoogleGenerativeAI(model=gemini_model, google_api_key=GOOGLE_API_KEY, temperature=0.8, system_prompt=system_prompt, response_validation=False)

            retriever = vectorstore.as_retriever()
            prompt = hub.pull("rlm/rag-prompt")

            def format_docs(docs):
                return "\n\n".join(doc.page_content for doc in docs)

            rag_chain = (
                {"context": retriever | format_docs, "question": RunnablePassthrough()}
                | prompt
                | llm
                | StrOutputParser()
            )

            # Interacción con el usuario
            pregunta_ = st.text_area("Haz una pregunta sobre los documentos:")
            pregunta = "Das respuestas en español,eres experto en analisis de datos,Siempre debes proporcionar respuestas claras,concisas y basadas en hechos." + pregunta_
            if pregunta:
                response = rag_chain.invoke(pregunta)
                st.markdown(response)

    if __name__ == "__main__":
        main()
else:
    st.write("Por favor, proporciona las API keys para continuar.")