import os
import streamlit as st
from tempfile import NamedTemporaryFile
from retriever import get_retrieval_chain
from preprocess import create_vectorstore
from preprocess import load_vectorstore

def rag_with_new_data(file_name):
    cwd = os.getcwd()
    temp_dir = os.path.join(cwd, "temp")
    os.makedirs(temp_dir, exist_ok=True)
    file_path = os.path.join(temp_dir, file_name.name)
    
    st.write("Creating a temporary file..")
    with open(file_path, "wb") as f:
        f.write(file_name.getbuffer())
        
    st.write("Temp file created...") 
    
    st.write("Wait while I create the embeddings..")
    vectorstore = create_vectorstore(f.name)
    st.write("OK the embeddings are ready.")

    retriever = "google/flan-t5-large"
    qa_chain = get_retrieval_chain(retriever, vectorstore)
    
    return qa_chain

def rag_with_saved_data(saved_db_name):

    vectorstore = load_vectorstore(saved_db_name)
    
    retriever = "google/flan-t5-large"
    qa_chain = get_retrieval_chain(retriever, vectorstore)
    
    return qa_chain
    
st.title("Ask me anything about the PDF!")

file_name =None
qa_chain = None

choice = st.sidebar.radio("select" , 
                          ["Existing Knowledge Base", "New Knowledge Base"])

if choice == "Existing Knowledge Base":
    
    selected_option= st.sidebar.selectbox("choose an existing knowledgebase", 
                                         ("choose an option", "underwriting"),
                                         index = 0,
                                         placeholder = "choose an option"
                                         )
    if selected_option == "underwriting":
        qa_chain = rag_with_saved_data(selected_option)
elif choice == "New Knowledge Base":
    file_name = st.sidebar.file_uploader("Upload a PDF file", type=['pdf'])
    if file_name:
        qa_chain = rag_with_new_data(file_name)
    

user_prompt = st.text_input('Ask your question here e.g. \"what is the total purchase price\"')

if user_prompt != '' and qa_chain:
    # with NamedTemporaryFile(dir='.', suffix='.pdf', mode='wb') as f:
    #     f.write(file_name.getbuffer())
    #     st.write("file written:", f.name)
    
    response = qa_chain(user_prompt)
    st.write(response)