import gradio as gr
import os
from dotenv import load_dotenv
from pathlib import Path
from functions import (
    reset_open_ai,
    fetch_collections,
    select_collection,
    reset_collection,
    handle_submit,
)
from models import (
    Icons,
    Message,
    Messages,
)

load_dotenv()

css_style = Path('./style.css').read_text()

with gr.Blocks(css=css_style) as demo:
    zot = gr.State(None)
    zot_collections = gr.State([])
    data = gr.State([])
    messages = gr.State(Messages([
        Message(
            Icons.INFO, "Please provide all the required OpenAI and Zotero information in the left panel."),
    ]))

    with gr.Row():
        gr.Markdown("""
        <h2>
          <img src='file/assets/zotero-logo.png' alt="Zotero Logo" style="width: 3.2rem; display: inline; margin-right: 10px;" />
          Zotero Q&A
        </h2>

        Fan Li ([@FanLi_RnD](https://twitter.com/FanLi_RnD)) - https://apex974.com/articles/literature-reviews-with-paper-qa-and-zotero
        
        This tool allows you to ask questions based on your Zotero library. It was built upon [Paper QA](https://github.com/whitead/paper-qa), [LangChain AI](https://github.com/hwchase17/langchain) and [pyZotero](https://github.com/urschrei/pyzotero).

        <b>Check out Aria, my new AI Research Assistant Plugin for Zotero: https://github.com/lifan0127/ai-research-assistant.</b>
        """)

    with gr.Row():
        with gr.Column(scale=1):
            openai_api_key = gr.Textbox(
                label="OpenAI API Key", type="password", value=os.getenv('OPENAI_API_KEY'))

            gr.HTML()

            zot_api_key = gr.Textbox(
                label="Zotero API Key", type="password", value=os.getenv('ZOTERO_API_KEY'))
            zot_library_type = gr.Radio(choices=[
                                        "User", "Group"], label="Zotero Library Type", value="User", elem_id="zotero-library-type")
            zot_library_id = gr.Textbox(
                label="Zotero User/Group ID", value=os.getenv('ZOTERO_LIBRARY_ID'))

            zot_citation_style = gr.Textbox(
                label="Zotero Citation Style",
                value='nature',
            )

            zot_selected_col = gr.Radio(
                [], label="Zotero Collection", elem_id="zotero-collection", visible=False)
            zot_msg = gr.HTML("""
            <div style="padding: 1rem; background-color: #fffbe7; font-size: 0.8rem;">
              <ul style="margin-bottom: 0;">
                <li>Click <a href="https://www.zotero.org/settings/keys/new" target="_blank">here</a> to create Zotero API key.</li>
                <li>To access your own library, select "User".</li>
                <li>To access a shared group, select "Group".</li>
                <li>Personal User ID can be found <a href="https://www.zotero.org/settings/keys" target="_blank">here</a>.</li>
                <li style="margin-bottom: 0;">Group ID is part of the group URL (e.g. 4952526).</li>
              </ul>
            </div>
            """, visible=True)
            zot_fetch_col_btn = gr.Button('Fetch Collections')

            gr.Error("Some Error Message")

        with gr.Column(scale=3):

            question = gr.Textbox(
                placeholder="You have to select a Zotero collection to proceed", label="Question", interactive=False)

            gr.HTML()

            with gr.Box():
                with gr.Accordion("Messages"):
                    msg_board = gr.HTML(messages.value)

                answer = gr.HTML(None, elem_id="answer")

    openai_api_key.change(reset_open_ai, inputs=[openai_api_key], outputs=[
                          answer], show_progress=False)

    zot_api_key.change(
        reset_collection,
        inputs=[messages],
        outputs=[zot_selected_col, zot_fetch_col_btn,
                 question, answer, messages, msg_board, zot_selected_col],
        show_progress=False,
    )
    zot_library_type.change(
        reset_collection,
        inputs=[messages],
        outputs=[zot_selected_col, zot_fetch_col_btn,
                 question, answer, messages, msg_board, zot_selected_col],
        show_progress=False,
    )
    zot_library_id.change(
        reset_collection,
        inputs=[messages],
        outputs=[zot_selected_col, zot_fetch_col_btn,
                 question, answer, messages, msg_board, zot_selected_col],
        show_progress=False
    )
    zot_citation_style.change(
        reset_collection,
        inputs=[messages],
        outputs=[zot_selected_col, zot_fetch_col_btn,
                 question, answer, messages, msg_board, zot_selected_col],
        show_progress=False
    )

    zot_fetch_col_btn.click(
        fn=fetch_collections,
        inputs=[openai_api_key, zot_library_id,
                zot_library_type, zot_api_key, messages],
        outputs=[zot, zot_collections, zot_selected_col,
                 zot_fetch_col_btn, zot_msg, messages, msg_board],
        show_progress=False,
    )
    zot_selected_col.change(
        fn=select_collection,
        inputs=[zot_selected_col, messages],
        outputs=[question, messages, msg_board, answer],
        show_progress=False
    )

    question.submit(
        fn=handle_submit,
        inputs=[zot, zot_selected_col, zot_collections,
                zot_citation_style, question, messages],
        outputs=[messages, msg_board, answer],
        show_progress=False
    )

demo.queue(concurrency_count=10, api_open=False)
demo.launch()