import gradio as gr
import openai
# from gpt_reader.pdf_reader import PaperReader
# from gpt_reader.prompt import BASE_POINTS
from Document_QA import QA
from Document_QA import create_embeddings
from Document_QA import Paper
from PyPDF2 import PdfReader

class GUI:
    def __init__(self):
        self.api_key = ""
        self.session = ""
        self.all_embedding =None
        self.tokens = 0
    #load pdf and create all embedings
    def pdf_init(self, api_key, pdf_path):
        openai.api_key = api_key
        pdf_reader = PdfReader(pdf_path.name)
        paper = Paper(pdf_reader)
        all_texts = paper.get_texts()
        self.all_embedding, self.tokens = create_embeddings(all_texts)
        print("全部文本消耗 {} tokens".format(self.tokens))

    def get_answer(self, question):
        qa = QA(self.all_embedding)
        answer,context = qa(question)
        return answer.strip()

with gr.Blocks() as demo:
    gr.Markdown(
        """
        # CHATGPT-PAPER-READER
        [点击此处以支付 $5 成为我们的会员](https://checkout.stripe.com/c/pay/cs_live_a1TwwqhUpsfstnbyiAvbMoXvMzoaII5vskE8tz1cIsMSYUt9hJvoHK2qOK#fidkdWxOYHwnPyd1blppbHNgWjA0TlZXUHNAck9nTWNdXVc1TDRxTXIzQGo9b383N11yfDBhMzBvZ0pAMlNURDBBVWpiMHJObkhkSUZQSktwaWZ9S1dqUzFRRDw0f1dSa0dAQmp%2FYk5TS2tQNTVHa1F1RlVvPCcpJ3VpbGtuQH11anZgYUxhJz8nZEBQZko9MWRMPDxEYUNOZkhIJ3gl)
        """)

    with gr.Tab("Upload PDF File"):
        pdf_input = gr.File(label="PDF File")
        api_input = gr.Textbox(label="OpenAI API Key")
        #result = gr.Textbox(label="PDF Summary")
        upload_button = gr.Button("Start Analyse")
    with gr.Tab("Ask question about your PDF"):
        question_input = gr.Textbox(label="Your Question", placeholder="Authors of this paper?")
        answer = gr.Textbox(label="Answer")
        ask_button = gr.Button("Ask")

    app = GUI()
    upload_button.click(fn=app.pdf_init, inputs=[api_input, pdf_input])
    ask_button.click(app.get_answer, inputs=question_input, outputs=answer)

if __name__ == "__main__":
    demo.title = "CHATGPT-PAPER-READER"
    demo.launch()  # add "share=True" to share CHATGPT-PAPER-READER app on Internet.