from response_parser import *
import gradio as gr
import os


username = os.getenv('USERNAME')
password = os.getenv('PASSWORD')

def initialization(state_dict: Dict) -> None:
    if not os.path.exists('cache'):
        os.mkdir('cache')
    if state_dict["bot_backend"] is None:
        state_dict["bot_backend"] = BotBackend()
        if 'OPENAI_API_KEY' in os.environ:
            del os.environ['OPENAI_API_KEY']


def get_bot_backend(state_dict: Dict) -> BotBackend:
    return state_dict["bot_backend"]


def switch_to_gpt4(state_dict: Dict, whether_switch: bool) -> None:
    bot_backend = get_bot_backend(state_dict)
    if whether_switch:
        bot_backend.update_gpt_model_choice("GPT-4")
    else:
        bot_backend.update_gpt_model_choice("GPT-3.5")


def add_text(state_dict: Dict, history: List, text: str) -> Tuple[List, Dict]:
    bot_backend = get_bot_backend(state_dict)
    bot_backend.add_text_message(user_text=text)

    history = history + [(text, None)]

    return history, gr.update(value="", interactive=False)


def add_file(state_dict: Dict, history: List, file) -> List:
    bot_backend = get_bot_backend(state_dict)
    path = file.name
    filename = os.path.basename(path)

    bot_msg = [f'📁[{filename}]', None]
    history.append(bot_msg)

    bot_backend.add_file_message(path=path, bot_msg=bot_msg)

    return history


def undo_upload_file(state_dict: Dict, history: List) -> Tuple[List, Dict]:
    bot_backend = get_bot_backend(state_dict)
    bot_msg = bot_backend.revoke_file()

    if bot_msg is None:
        return history, gr.Button.update(interactive=False)

    else:
        assert history[-1] == bot_msg
        del history[-1]
        if bot_backend.revocable_files:
            return history, gr.Button.update(interactive=True)
        else:
            return history, gr.Button.update(interactive=False)


def refresh_file_display(state_dict: Dict) -> List[str]:
    bot_backend = get_bot_backend(state_dict)
    work_dir = bot_backend.jupyter_work_dir
    filenames = os.listdir(work_dir)
    paths = []
    for filename in filenames:
        paths.append(
            os.path.join(work_dir, filename)
        )
    return paths


def restart_ui(history: List) -> Tuple[List, Dict, Dict, Dict, Dict]:
    history.clear()
    return (
        history,
        gr.Textbox.update(value="", interactive=False),
        gr.Button.update(interactive=False),
        gr.Button.update(interactive=False),
        gr.Button.update(interactive=False)
    )


def restart_bot_backend(state_dict: Dict) -> None:
    bot_backend = get_bot_backend(state_dict)
    bot_backend.restart()


def bot(state_dict: Dict, history: List) -> List:
    bot_backend = get_bot_backend(state_dict)

    while bot_backend.finish_reason in ('new_input', 'function_call'):
        if history[-1][0] is None:
            history.append(
                [None, ""]
            )
        else:
            history[-1][1] = ""

        response = chat_completion(bot_backend=bot_backend)
        for chunk in response:
            history, weather_exit = parse_response(
                chunk=chunk,
                history=history,
                bot_backend=bot_backend
            )
            yield history
            if weather_exit:
                exit(-1)

    yield history


if __name__ == '__main__':
    config = get_config()
    custom_css = """
    .gradio-container {
        background-color: white ;
        margin: 0 !important ;
        padding: 0 !important ;
        space : 0 
    }

    #mainDiv {
    width :100%;
    border : none;
    height: 100vh ;
    }

    #chatbot_div{
    padding : 10px
    }

    #chatbot {
    border-color: #a6a6a6 ;
    background-color: white ;
    border-radius: 5px;
    }

    #sidebar {
    background-color: #f2f2f2;
    padding : 5px ; 
    }

    #files {
   
    height : 60% ;
    color : #f2f2f2  ;
    background-color : #f2f2f2 ;
    border : none
    }

    #gpt4_button {
    border-radius:  0px ;
    background-color : #0CAFFF ;
    padding : 10px ;
    }

    #gpt4_button:hover {
    background-color : #0ca6ff ;
    }

    #textbox {
    border-color : #a6a6a6 ;
    background-color : white ;
    }

    #textbox:hover {
    border-color : black ;
    }

    #upload_button {
    background-color : #b3b3b3;
    }

    #upload_button:hover {
    background-color : #8c8c8c ;
    box-shadow: 0 12px 16px 0 rgba(0,0,0,0.24),0 17px 50px 0 rgba(0,0,0,0.19);
    }

    /* width */
    ::-webkit-scrollbar {
      width: 4px;
    }

    /* Track */
    ::-webkit-scrollbar-track {
      background: #f1f1f1;
    }

    /* Handle */
    ::-webkit-scrollbar-thumb {
      background: #0CAFFF;
    }

    /* Handle on hover */
    ::-webkit-scrollbar-thumb:hover {
      background: #555;
    }

    footer{display:none !important}

    """

    javascript_code = """
    function get_browser_height() {
    return window.innerHeight;
    }
    """
    with gr.Blocks(theme=gr.themes.Base(),css= custom_css , title='DMO-GPT-Interpreter') as block:
        """
        Reference: https://www.gradio.app/guides/creating-a-chatbot-fast
        """
        # UI components
        state = gr.State(value={"bot_backend": None})
        with gr.Row( elem_id="mainDiv"):
            with gr.Column( elem_id="sidebar" , scale=0.20):
                sidebar_header = gr.HTML("<h1 style='color:black; font-weight:bold; text-align:center; margin-top : 10px ;'>DMO-GPT-Interpreter</h1>")
                hr_linr = gr.HTML("<hr/>")
                file_section = gr.HTML("<h2 style='color:gray; font-weight:bold; text-align:start; margin-top : 5px ; font-size: 18px ; margin-bottom: -3px'>Files</h2>")
                file_output = gr.Files(elem_id="files", show_label=False )
                setting_section = gr.HTML("<h2 style='color:gray; font-weight:bold; text-align:start; margin-top : 5px ; font-size: 18px ; margin-bottom: -3px'>Settings</h2>")
                check_box = gr.Checkbox(label="Use with GPT-4 ✨", interactive=config['model']['GPT-4']['available'] , elem_id="gpt4_button" , scale=2)
                check_box.change(fn=switch_to_gpt4, inputs=[state, check_box])
            with gr.Column( elem_id="chatbot_div"):
                chatbot = gr.Chatbot([], elem_id="chatbot", height=600 , label="Welcome to DMO-GPT-Interpreter!" , show_share_button=False)
                with gr.Row():
                    with gr.Column(scale=0.85  ):
                        text_box = gr.Textbox(
                            show_label=False,
                            placeholder="Enter text and press enter, or upload a file",
                            container=False,
                            elem_id='textbox'
                    )
                    with gr.Column(scale=0.15, min_width=0):
                        file_upload_button = gr.UploadButton("📁 Upload log files", file_types=['file']  , elem_id="upload_button")
                
        # with gr.Tab("Chat"):
        #     chatbot = gr.Chatbot([], elem_id="chatbot", label="Local Code Interpreter", height=500)
        #     with gr.Row():
        #         with gr.Column(scale=0.85  ):
        #             text_box = gr.Textbox(
        #                 show_label=False,
        #                 placeholder="Enter text and press enter, or upload a file",
        #                 container=False
        #                 
        #             )
        #         with gr.Column(scale=0.15, min_width=0):
        #             file_upload_button = gr.UploadButton("📁", file_types=['file'])
            # with gr.Row(equal_height=True):
            #     with gr.Column(scale=0.7):
            #         check_box = gr.Checkbox(label="Use GPT-4", interactive=config['model']['GPT-4']['available'])
            #         check_box.change(fn=switch_to_gpt4, inputs=[state, check_box])
            #     with gr.Column(scale=0.15, min_width=0):
            #         restart_button = gr.Button(value='🔄 Restart')
            #     with gr.Column(scale=0.15, min_width=0):
            #         undo_file_button = gr.Button(value="↩️Undo upload file", interactive=False)
        # with gr.Tab("Files"):
        #     file_output = gr.Files()

        # Components function binding
        txt_msg = text_box.submit(add_text, [state, chatbot, text_box], [chatbot, text_box], queue=False).then(
            bot, [state, chatbot], chatbot
        )
        txt_msg.then(fn=refresh_file_display, inputs=[state], outputs=[file_output])
        txt_msg.then(lambda: gr.update(interactive=True), None, [text_box], queue=False)
        txt_msg.then(lambda: gr.Button.update(interactive=False), None,  queue=False)

        file_msg = file_upload_button.upload(
            add_file, [state, chatbot, file_upload_button], [chatbot], queue=False
        ).then(
            bot, [state, chatbot], chatbot
        )
        file_msg.then(lambda: gr.Button.update(interactive=True), None, queue=False)
        file_msg.then(fn=refresh_file_display, inputs=[state], outputs=[file_output])

        block.load(fn=initialization, inputs=[state] , _js=javascript_code)

    block.queue()
    block.launch(inbrowser=True , auth=(username,password) , auth_message= "Login to continue with DMO-GPT-Interpreter" , favicon_path="ai.png" )