import gradio as gr
from llm import ChatGLM, OpenAI3, OpenAI4
from pathlib import Path

block_css = """
.importantButton {
    background: linear-gradient(45deg, #7e0570,#5d1c99, #6e00ff) !important;
    border: none !important;
}
.importantButton:hover {
    background: linear-gradient(45deg, #ff00e0,#8500ff, #6e00ff) !important;
    border: none !important;
}"""

webui_title = "📚📚📚📚📚📚📚📚📚📚📚📚* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *    ☘️ * * *智海文心* * *    ☘️ * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *📚📚📚📚📚📚📚📚📚📚📚📚"

generate_prompt = 'please generate {generate_number} difficult multiple-choice questions and attach the answer to each ' \
                  'question according the paragraph below. Each question should have four choices and only one of them ' \
                  'is the right answer.' \

correct_prompt = 'the following paragraph is a composition of a student. ' \
                 'please first judge the writing level of the composition, ' \
                 'then point out the all the problems with examples, ' \
                 'and finally output the complete revised version.'

llm_name_dict = {'chatgpt-3.5': 'OpenAI3', 'chatgpt-4': 'OpenAI4', 'chatglm—6b': 'ChatGLM'}


def files_list_to_texts(files_list, glob="**/[!.]*"):
    texts = []
    for i in files_list:
        if i.name.split('.')[-1] in ['txt']:
            with open(i.name, encoding='utf-8') as f:
                text = f.readlines()
                text = [i for i in text]
                texts.append(''.join(text))
    return texts


def function_select(mode):
    if mode == "问题生成":
        return gr.update(visible=True), gr.update(visible=False)
    else:
        return gr.update(visible=False), gr.update(visible=True)


def llm_change(name):
    llm = eval(eval('llm_name_dict[name]'))()
    return llm


def text_generate(chatbot, text, generate_number, llm):
    prompt = eval('f"' + generate_prompt + '"') + '\n\n' + text
    answer = llm(prompt)
    chatbot = chatbot + [[text, answer]]
    return chatbot


def files_generate(chatbot, files_list, generate_number, llm):
    try:
        texts = files_list_to_texts(files_list)
        for text in texts:
            prompt = eval('f"' + generate_prompt + '"') + '\n\n' + text
            answer = llm(prompt)
            chatbot = chatbot + [[text, answer]]
    except:
        chatbot = chatbot + [[None, f"所选文件夹中的文件添加失败，请确保文件夹中含有txt类型文件"]]
    return chatbot


def text_correct(chatbot, text, llm):
    prompt = correct_prompt + '\n\n' + text
    answer = llm(prompt)
    chatbot = chatbot + [[text, answer]]
    return chatbot


def files_correct(chatbot, files_list, llm):
    try:
        texts = files_list_to_texts(files_list)
        for text in texts:
            prompt = correct_prompt + '\n\n' + text
            answer = llm(prompt)
            chatbot = chatbot + [[text, answer]]
    except:
        chatbot = chatbot + [[None, f"所选文件夹中的文件添加失败，请确保文件夹中含有txt类型文件"]]
    return chatbot


def clear_screen(chatbot):

    return [[None, None]]


with gr.Blocks(css=block_css) as demo:
    gr.Markdown('\n\n\n\n')
    gr.Markdown(webui_title)
    gr.Markdown('\n\n\n\n')

    llm = gr.State('')
    model_mode = gr.Radio(['chatglm—6b', "chatgpt-3.5", "chatgpt-4"], label="请选择驱动模型")
    model_mode.change(fn=llm_change, inputs=[model_mode], outputs=[llm])

    fun_mode = gr.Radio(["问题生成", "作文批改"], label="请选择功能模式")
    qg = gr.Row(visible=False)
    aa = gr.Row(visible=False)
    fun_mode.change(fn=function_select, inputs=[fun_mode], outputs=[qg, aa])

    with qg:

        with gr.Column(scale=10):
            chatbot = gr.Chatbot([[None, None]],
                                 elem_id="chat-box",
                                 show_label=False).style(height=800)
            clear_button = gr.Button(value="清屏")
            clear_button.click(fn=clear_screen, inputs=[chatbot], outputs=[chatbot])

        with gr.Column(scale=10):

            with gr.Tab('生成配置'):
                generate_number = gr.Slider(1,
                                            5,
                                            value=3,
                                            step=1,
                                            label="请设定单篇文章需要生成的问题数量",
                                            interactive=True)

                gr.Markdown(f'单篇生成')
                texts = gr.Textbox(show_label=False, placeholder="文本内容", lines=12).style(container=False)
                text_button = gr.Button(value="生成问题")
                text_button.click(fn=text_generate, inputs=[chatbot, texts, generate_number, llm], outputs=[chatbot])

                gr.Markdown(f'批量生成')
                folder_address = gr.File(label="添加文件",
                                file_types=['.txt', '.md', '.docx', '.pdf'],
                                file_count="multiple",
                                show_label=False
                                )
                file_button = gr.Button(value="生成问题")
                file_button.click(fn=files_generate, inputs=[chatbot, folder_address, generate_number, llm],
                                  outputs=[chatbot])


    with aa:

        with gr.Column(scale=10):
            chatbot = gr.Chatbot([[None, None]],
                                 elem_id="chat-box",
                                 show_label=False).style(height=800)
            clear_button = gr.Button(value="清屏")
            clear_button.click(fn=clear_screen, inputs=[chatbot], outputs=[chatbot])

        with gr.Column(scale=10):

            with gr.Tab('批改配置'):

                gr.Markdown(f'单篇批改')
                texts = gr.Textbox(show_label=False, placeholder="文本内容", lines=16).style(container=False)
                text_button = gr.Button(value="批改")
                text_button.click(fn=text_correct, inputs=[chatbot, texts, llm], outputs=[chatbot])

                gr.Markdown(f'批量批改')
                folder_address = gr.File(label="添加文件",
                                         file_types=['.txt', '.md', '.docx', '.pdf'],
                                         file_count="multiple",
                                         show_label=False
                                         )
                file_button = gr.Button(value="批改")
                file_button.click(fn=files_correct, inputs=[chatbot, folder_address, llm], outputs=[chatbot])


demo.queue(concurrency_count=5).launch(show_api=False)