import gradio as gr
import base64
import os 
api_key = os.getenv('API_KEY')

def predict(input, file_input):
    print("input:", input)
    print("file_input:", file_input.name)
    from gradio_client import Client
    
    client = Client(api_key)
    extract_result = client.predict(
        input,
        file_input.name,
        fn_index=1
    )
    if extract_result:
        print(extract_result)
        return extract_result
    else:
        return "Too many user, please wait a monument!"


def view_pdf(pdf_file):
    with open(pdf_file.name, 'rb') as f:
        pdf_data = f.read()
    # print("pdf_file", pdf_file)
    # pdf_data = pdf_file
    b64_data = base64.b64encode(pdf_data).decode('utf-8')
    # print("b64_data", b64_data)
    return f"<embed src='data:application/pdf;base64,{b64_data}' type='application/pdf' width='100%' height='700px' />"


en_1 = ["""could you please help me extract the information of 'title'/'journal'/'year'/'author'/'institution'/'email' from the previous content in a markdown table format?
If any of this information was not available in the paper, please replaced it with the string `""`, If the property contains multiple entities, please use a list to contain.
"""]

en_2 = ["""could you please help me extract the information of 'title'/'journal'/'year'/'author'/'institution'/'email' from the previous content in a json format?
If any of this information was not available in the paper, please replaced it with the string `""`, If the property contains multiple entities, please use a list to contain.
"""]

examples = [en_1, en_2]

with gr.Blocks(title="ChatPaperGPT") as demo:
    gr.Markdown(
        '''<p align="center" width="100%">
        <img src="https://big-cheng.com/img/pdf.png" alt="pdf-logo" width="50"/>
        <p>

        <h1 align="center"> Paper Extract GPT </h1>
        <p> How to use:
        <br> <strong>#1</strong>: Upload your pdf.
        <br> <strong>#2</strong>: Click the View PDF button to view it.
        <br> <strong>#3</strong>: Enter your extraction prompt in the input box (of course, you can click example to test).
        <br> <strong>#4</strong>: Click Generate to extract, and the extracted information will be displayed in markdown form.
        </p>
        '''
    )
    with gr.Row():
        with gr.Column():
            gr.Markdown('## Upload PDF')
            file_input = gr.File(type="filepath")
            viewer_button = gr.Button("View PDF")
            file_out = gr.HTML()
        with gr.Column():
            with gr.Row():
                model_input = gr.Textbox(lines=7, placeholder='Input prompt about extract information from paper',
                                         label='Input')
            with gr.Row():
                gen = gr.Button("Generate")
                clr = gr.Button("Clear")
            example = gr.Examples(examples=examples, inputs=model_input)

            with gr.Row():
                outputs = gr.Markdown(label='Output', show_label=True, value="""| Title                                       | Journal            | Year | Author                                        | Institution                                           | Email                 |
|---------------------------------------------|--------------------|------|-----------------------------------------------|-------------------------------------------------------|-----------------------|
| Paleomagnetic Study of Deccan Traps from Jabalpur to Amarkantak, Central India | J. Geomag. Geoelectr. | 1973 | R. K. VERMA, G. PULLAIAH, G.R. ANJANEYULU, P. K. MALLIK | National Geophysical Research Institute, Hyderabad, and Indian School o f Mines, Dhanbad | "" |
""")

    inputs = [model_input, file_input]
    gen.click(fn=predict, inputs=inputs, outputs=outputs)
    clr.click(fn=lambda value: [gr.update(value=""), gr.update(value="")], inputs=clr,
              outputs=[model_input, outputs])

    viewer_button.click(view_pdf, inputs=file_input, outputs=file_out)
    # parser_button.click(extract_text, inputs=file_input, outputs=[xml_out, md_out, rich_md_out])

demo.launch()