import subprocess
import os
import streamlit as st
import os.path
import pathlib
import sqlite3 
import re

def upload_pdf():
    st.set_page_config(page_title="Upload PDF", page_icon=":pdf:", layout="centered", initial_sidebar_state="expanded")
    file_upload = st.file_uploader("Upload your PDF document", type=["pdf"])
    if file_upload is not None:
        st.success("PDF uploaded successfully!")
        if st.button('Submit'):
            process_pdf(file_upload)

def process_pdf(file):
    # This function will call the specific command on the PDF file
    st.write("Processing PDF...")
    st.write("Name of the pdf: ", file.name)
    abs_path = os.path.abspath(file.name)
    with open(abs_path, "wb") as f:
        f.write(file.getvalue())
    openai_api_key = "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"
    if openai_api_key is None:
        st.error("Did not find openai_api_key, please add an environment variable `OPENAI_API_KEY` which contains it")
    else:
        openai_api = f"export OPENAI_API_KEY={openai_api_key}"
        command = f"dr-doc-search --train -i '{abs_path}'"
        result = subprocess.run(command, shell=True)
        if result.returncode == 0:
            st.success("Process completed!")
            create_or_connect_db()
            insert_pdf_data(file, abs_path)
            st.success("Data stored successfully!")            
        else:
            st.error("An error occurred while processing the PDF.")
            #st.write("stderr: ", result.stderr.decode())       
  
def create_or_connect_db():
    conn = sqlite3.connect('pdf_docs.db')
    c = conn.cursor()
    c.execute('''CREATE TABLE IF NOT EXISTS pdf_docs
    (id INTEGER PRIMARY KEY, file_name TEXT, abs_path TEXT)''')
    conn.commit()
    return conn

def insert_pdf_data(file, abs_path):
    conn = create_or_connect_db()
    c = conn.cursor()
    c.execute("INSERT INTO pdf_docs (file_name, abs_path) VALUES (?, ?)", (file.name, abs_path))
    conn.commit()
    st.success("PDF data inserted into the database!")

def get_all_pdf_data():
    conn = create_or_connect_db()
    c = conn.cursor()
    c.execute("SELECT file_name FROM pdf_docs")
    pdf_data = c.fetchall()
    pdf_file_names = [x[0] for x in pdf_data]
    return pdf_file_names

def ask():
    pdf_file_names = get_all_pdf_data()
    pdf_file_name = st.selectbox("Select a PDF file", pdf_file_names)
    if pdf_file_name:
        st.write("You selected: ", pdf_file_name)
        st.title("Ask a question about a PDF")
        question = st.text_area(
                    "Use the example below or input your own text in English",
                    value="test",
                    max_chars=200,
                    height=100
                    )
        if st.button("Submit"):
            with st.spinner(text="In progress"):
                command = f"dr-doc-search -q '{question}' -i '{pdf_file_name}' -v"
                result = subprocess.run(command, shell=True, stdout=subprocess.PIPE)
                report_text = result.stdout.decode()
                report_text="Question: "+"Question: ".join(report_text.split("Question: ")[1:])
                #report_text = re.search(r'(?<=Answer:).*', report_text).group(0)
                #report_text = re.search(r'(?<=Question:).*', report_text).group(0)
                st.write(report_text + re.search(r'(?<=Answer: ).*', report_text).group(0))


def main():
    st.set_page_config(page_title="Welcome", page_icon=":guardsman:", layout="wide")    
    st.sidebar.title("Options")
    menu = ["Ask", "Upload"]
    choice = st.sidebar.selectbox("Select an option", menu)

    if choice == "Upload":
        st.title("Index a new pdf document")
        file = st.file_uploader("Upload a PDF document", type=["pdf"])
        if file is not None:
            process_pdf(file)
    elif choice == "Ask":
        ask()
    else:
        st.warning("Please select a valid option from the sidebar.")

if __name__=='__main__':
    main()