import gradio as gr
import openai


def get_answer_chatgpt(chat_history, message, system_message=None):
    messages = []
    if system_message:
        TEMPLATE_SYSTEM = {"role": "system", "content": system_message}
        message.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="gpt-3.5-turbo",
        messages=messages
    )

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



with gr.Blocks() as demo:
    system_prompt = 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")
    clear = gr.Button("Clear")
    send = gr.Button('Send')

    def respond(chat_history, message):
        response = get_answer_chatgpt(chat_history, message)
        response = response["completion"]

        return chat_history + [[message, response]]

    message.submit(respond, [chatbot, message], chatbot)

    clear.click(lambda: None, None, chatbot, queue=False)
    run.click(respond, [chatbot, message], chatbot)

    demo.launch(debug=True)