import os 
import streamlit as st
from langchain.chains import RetrievalQA
from langchain.llms import OpenAI
from langchain.document_loaders import TextLoader
from langchain.document_loaders import PyPDFLoader
from langchain.indexes import VectorstoreIndexCreator
from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.document_loaders import PyPDFDirectoryLoader
import streamlit as st
import os
import openai

from langchain.chains.question_answering import load_qa_chain

# Create a title for the app
st.title("ChromaDB Multiple PDFs")

st.markdown("**OpenAI API key**")
key = st.text_input("Paste Your API key here")
st.write(key)

if key:

    os.environ["OPENAI_API_KEY"] = key
    # Create a sidebar for selecting the number of files to upload
    st.sidebar.header("Number of files")
    num_files = st.sidebar.number_input("How many PDF files do you want to upload?", min_value=1, max_value=5, value=1)

    # Create a list to store the uploaded files
    uploaded_files = []

    # Loop through the number of files and create file uploaders
    for i in range(num_files):
        uploaded_file = st.file_uploader(f"Choose a PDF file {i+1}", type="pdf")
        # If a file is uploaded, append it to the list
        if uploaded_file is not None:
            uploaded_files.append(uploaded_file)

    # Check if any file is uploaded
    if len(uploaded_files) > 0:
        # Check if the books folder exists, if not create it
        if not os.path.exists("books"):
            os.mkdir("books")
        # Loop through the uploaded files and save them to the books folder
        for i, file in enumerate(uploaded_files):
            # Create a file name with the index and the original name
            file_name = f"{i}_{file.name}"
            # Open the file in binary mode and write its contents
            with open(os.path.join("books", file_name), "wb") as f:
                f.write(file.getbuffer())
        # Display a success message
        st.success(f"Successfully uploaded {len(uploaded_files)} PDF files to the books folder.")







    # load document
    loader = PyPDFDirectoryLoader("books/")
    documents = loader.load()

### For multiple documents 
# loaders = [....]
# documents = []
# for loader in loaders:
#     documents.extend(loader.load())

    chain = load_qa_chain(llm=OpenAI(), chain_type="map_reduce")
    query = st.text_input("Write your query")
    if query:
        lang_agent_run = chain.run(input_documents=documents, question=query)
        st.write("Lang chain agent answer: ",lang_agent_run )