import gradio as gr
import os
import sys
import argparse
from utils import *
from presets import *
from chat_func import *

my_api_key = os.environ.get('my_api_key')

if my_api_key == "empty":
    print("Please give a api key!")
    sys.exit(1)

gr.Chatbot.postprocess = postprocess


with open("custom.css", "r", encoding="utf-8") as f:
    customCSS = f.read()

with gr.Blocks(
    css=customCSS,
    # theme=gr.themes.Soft(
    #     primary_hue=gr.themes.Color(
    #         c50="#7b68ee"),
    #     secondary_hue=gr.themes.Color(
    #         c50="#576b95"),
    #     neutral_hue=gr.themes.Color(
    #         name="gray"),
        radius_size=gr.themes.sizes.radius_sm,
    ).set(
        button_primary_background_fill="#708090",
        button_primary_background_fill_dark="#778899",
        button_primary_background_fill_hover="#778899",
        button_primary_border_color="#ffffff",
        button_primary_border_color_dark="#ffffff",
        button_primary_text_color="#FFFFFF",
        button_primary_text_color_dark="#FFFFFF",
        button_secondary_background_fill="#F2F2F2",
        button_secondary_background_fill_dark="#2B2B2B",
        button_secondary_text_color="#393939",
        button_secondary_text_color_dark="#FFFFFF",
        block_title_text_color="#000000",
        block_title_background_fill="#000000",
        input_background_fill="#F6F6F6",
    ),
) as demo:
    
    history = gr.State([])
    token_count = gr.State([])
    promptTemplates = gr.State(load_template(get_template_names(plain=True)[0], mode=2))
    user_api_key = gr.State(my_api_key)
    TRUECOMSTANT = gr.State(True)
    FALSECONSTANT = gr.State(False)
    gr.Markdown(title)

    with gr.Accordion("Build by [45度科研人](WeChat Public Accounts)", open=False):
        gr.Markdown(description)
        
    with gr.Row(scale=1).style(equal_height=True):
        with gr.Column(scale=5):
            with gr.Column():
                chatbot = gr.Chatbot().style(colormap=("blue", "green"))
                user_input = gr.Textbox(show_label=False, placeholder="Enter text and press submit", visible=True).style(container=False)
                submitBtn = gr.Button("Submit", variant="primary").style(container=False)
                emptyBtn = gr.Button("Restart Conversation")
                status_display = gr.Markdown("")

        with gr.Column():
            with gr.Column(min_width=50):
                with gr.Tab(label="ChatGPT"):
                    with gr.Column():
                        with gr.Row():
                            keyTxt = gr.Textbox(show_label=False, placeholder=f"You can input your own openAI API-key",value=hide_middle_chars(my_api_key),visible=True, type="password",  label="API-Key")
                            # keyTxt = gr.Textbox(show_label=False, placeholder=f"You can input your own openAI API-key",value=my_api_key,visible=True, type="password",  label="API-Key")
                            systemPromptTxt = gr.Textbox(show_label=True,placeholder=f"Set a custom insruction for the chatbot: You are a helpful assistant.",label="Custom prompt",value=initial_prompt,lines=10,).style(container=False)

                        with gr.Row():
                            templateSelectDropdown = gr.Dropdown(label="load from template",choices=load_template(get_template_names(plain=True)[0], mode=1),
                                multiselect=False,value=load_template(get_template_names(plain=True)[0], mode=1)[0],).style(container=False)                
                
                with gr.Tab(label="Settings"):
                    with gr.Column():
                        with gr.Row():
                            with gr.Column(scale=3):
                                saveFileName = gr.Textbox(show_label=True, placeholder=f"output file name...",label='Save conversation history', value="").style(container=False)
                            with gr.Column(scale=1):
                                exportMarkdownBtn = gr.Button("Save")
                        with gr.Row():
                            with gr.Column(scale=1):
                                downloadFile = gr.File(interactive=False)
    gr.Markdown("""
                                <div align=center>you can follow the WeChat public account [45度科研人] and leave me a message! 
                                <div align=center><img width = '200' height ='200' src ="https://dunazo.oss-cn-beijing.aliyuncs.com/blog/wechat-simple.png"/>  <img width = '170' height ='190' src ="https://dunazo.oss-cn-beijing.aliyuncs.com/blog/shoukuanma222.png"/>
                                </div>""")
    keyTxt.submit(submit_key, keyTxt, [user_api_key, status_display])
    keyTxt.change(submit_key, keyTxt, [user_api_key, status_display])
    # Chatbot
    user_input.submit(predict,[user_api_key,systemPromptTxt,history,user_input,chatbot,token_count,],[chatbot, history, status_display, token_count],show_progress=True)
    user_input.submit(reset_textbox, [], [user_input])

    submitBtn.click(predict,[user_api_key,systemPromptTxt,history,user_input,chatbot,token_count,],[chatbot, history, status_display, token_count],show_progress=True)
    submitBtn.click(reset_textbox, [], [user_input])

    emptyBtn.click(reset_state,outputs=[chatbot, history, token_count, status_display],show_progress=True,)

    templateSelectDropdown.change(get_template_content,[promptTemplates, templateSelectDropdown, systemPromptTxt],[systemPromptTxt],show_progress=True,)
    exportMarkdownBtn.click(export_markdown,[saveFileName, systemPromptTxt, history, chatbot],downloadFile,show_progress=True,)
    downloadFile.change(load_chat_history,[downloadFile, systemPromptTxt, history, chatbot],[saveFileName, systemPromptTxt, history, chatbot],)

# demo.title = "Sydney-AI 2.0"

if __name__ == "__main__":
    demo.queue().launch(debug=False,show_api=False)