# Import necessary packages
from llama_index import GPTSimpleVectorIndex, download_loader
from pathlib import Path
import os
import json
import gradio as gr 
import tempfile

def construct_index(file_path):
    PDFReader = download_loader("PDFReader")

    loader = PDFReader()
    documents = loader.load_data(file=Path(file_path))

    # Construct a simple vector index
    index = GPTSimpleVectorIndex.from_documents(documents)

    # Save your index to a index.json file
    index.save_to_disk('index.json')

    return index

def qabot(file, input_text):
    # Check if index already exists
    if not os.path.exists('index.json'):
        # If index does not exist, create index from file
        index = construct_index(file.name)
    else:
        # If index exists, load index from file
        index = GPTSimpleVectorIndex.load_from_disk('index.json')

    # Query the index with the user's input text
    response = index.query(input_text, response_mode="compact")
    return response.response

# Add input component for file upload
file_upload = gr.inputs.File(label="Upload PDF file")

# Change the input components of the function to the file upload component and a text box for user input
iface = gr.Interface(fn=qabot, inputs=[file_upload, gr.inputs.Textbox(lines=7, label='Enter your query')], outputs="text", title="Custom-trained QA Application")

# Add a separate interface to update the index
def update_index(file):
    # Save the uploaded file to a temporary file
    with tempfile.NamedTemporaryFile(delete=False) as temp_file:
        temp_file.write(file.read())
        temp_file_path = temp_file.name

    # Construct the index from the temporary file
    index = construct_index(temp_file_path)

    # Remove the temporary file
    os.remove(temp_file_path)

    # Update the index file
    index.save_to_disk('index.json')

    return "Index generated from uploaded file: {}".format(file.name)

update_index_interface = gr.Interface(update_index, inputs=file_upload, outputs="text", title="Update Index")

# Launch both interfaces
iface.launch()
update_index_interface.launch()