import gradio as gr
import openai


def get_answer_chatgpt(chat_history, message, system_message=None, model="gpt-4", max_length=1000):
    messages = []
    if model is None:
        model = "gpt-4"

    if system_message or system_message != "":
        TEMPLATE_SYSTEM = {"role": "system", "content": system_message}
        messages.append(TEMPLATE_SYSTEM)

    TEMPLATE_USER = {"role": "user"}
    TEMPLATE_BOT = {"role": "assistant"}

    # add historical conversation
    for pair in chat_history:
        TEMPLATE_USER["content"] = pair[0]
        TEMPLATE_BOT["content"] = pair[1]
        messages.append(TEMPLATE_USER)
        messages.append(TEMPLATE_BOT)
    # add latest input
    TEMPLATE_USER["content"] = message
    messages.append(TEMPLATE_USER)
    response = openai.ChatCompletion.create(
        model=model,
        messages=messages,
        max_tokens=max_length
    )

    return {
        "completion": response["choices"][0]["message"]["content"],
        "usage_prompt_token": response["usage"]["prompt_tokens"],
        "usage_completion_token": response["usage"]["completion_tokens"]
    }


def reset_textbox():
    return gr.update(value='')


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column(scale=5):
            clear = gr.Button("Clear all converstation")
        with gr.Column(scale=5):
            with gr.Accordion("Parameters", open=False):
                model = gr.Dropdown(["gpt-4", "gpt-3.5-turbo"], label="Model", value="gpt-4")
                max_length = gr.Slider(0, 10000, value=1000, step=1, label="Maximum length token", interactive=True)

    system_message = gr.Textbox(placeholder="You are helpful assistence", lines=5, label="System Message")
    chatbot = gr.Chatbot()
    message = gr.Textbox(placeholder="Halo kak, aku mau bertanya", label="Chat Here")
    with gr.Row():
        with gr.Column(scale=5):
            send = gr.Button("Send")
        with gr.Column(scale=5):
            retry = gr.Button("Retry")

    def respond(chat_history, message, system_message, model, max_length=1000):
        response = get_answer_chatgpt(chat_history, message, system_message, model, max_length)
        response = response["completion"]

        return chat_history + [[message, response]]

    message.submit(respond, [chatbot, message, system_message, model, max_length], chatbot)

    clear.click(lambda: None, None, chatbot, queue=False)
    send.click(respond, [chatbot, message, system_message, model, max_length], chatbot)
    send.click(reset_textbox, [], [message])
    message.submit(reset_textbox, [], [message])

    retry.click(respond, [chatbot, message, system_message, model, max_length], chatbot)
    retry.click(reset_textbox, [], [message])

    demo.launch(debug=True)