import streamlit as st
import os

from PyPDF2 import PdfReader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings 
from langchain.vectorstores import FAISS
from langchain.chat_models import ChatOpenAI
from langchain.chains.question_answering import load_qa_chain

st.set_page_config('preguntaDOC')
st.header("Pregunta a tu PDF")
OPENAI_API_KEY = st.text_input('OpenAI API Key', type='password')
pdf_files = st.file_uploader("Carga tus documentos", type="pdf", accept_multiple_files=True, on_change=st.cache_resource.clear)

@st.cache_resource 
def create_embeddings(pdfs):
    text = ""
    for pdf in pdfs:
        pdf_reader = PdfReader(pdf)
        for page in pdf_reader.pages:
            text += page.extract_text()

    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=800,
        chunk_overlap=100,
        length_function=len
        )        
    chunks = text_splitter.split_text(text)

    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
    knowledge_base = FAISS.from_texts(chunks, embeddings)

    return knowledge_base

# Inicializa el estado de sesión para la pregunta del usuario y el historial
if 'user_question' not in st.session_state:
    st.session_state['user_question'] = ''
if 'history' not in st.session_state:
    st.session_state['history'] = []

if pdf_files:
    knowledge_base = create_embeddings(pdf_files)
    user_question = st.text_input("Haz una pregunta sobre tus PDFs:", value=st.session_state['user_question'])

    if st.button("Preguntar"):
        os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY
        docs = knowledge_base.similarity_search(user_question, 3)
        llm = ChatOpenAI(model_name='gpt-3.5-turbo')
        chain = load_qa_chain(llm, chain_type="stuff")
        respuesta = chain.run(input_documents=docs, question=user_question)

        # Añade la pregunta y la respuesta al historial
        st.session_state['history'].append((user_question, respuesta))
        
        # Limpiar el cuadro de texto después de mostrar la respuesta
        st.session_state['user_question'] = ''

        # Forzar la recarga de la página para actualizar el cuadro de texto
        st.experimental_rerun()

    # Mostrar el historial de preguntas y respuestas
    st.write("### Historial de preguntas y respuestas")
    for question, answer in st.session_state['history']:
        st.write(f"**Pregunta:** {question}")
        st.write(f"**Respuesta:** {answer}")