import os

import gradio as gr
import openai

openai.api_key = os.environ.get('OPENAI_API_KEY')
with open("system-prompt.txt", encoding="utf-8") as f:
    system_prompt = f.read()


def query(message: str, chat_history: list) -> str:
    global system_prompt
    gpt_messages = [
        {"role": "system", "content": system_prompt}
    ]
    for user_message, bot_message in chat_history:
        gpt_messages.extend([
            {"role": "user", "content": user_message},
            {"role": "assistant", "content": bot_message},
        ])
    gpt_messages.append({"role": "user", "content": message})

    rez = openai.ChatCompletion.create(
        model="gpt-4",
        messages=gpt_messages,
    )
    return rez['choices'][0]['message']['content']


with gr.Blocks() as demo:
    chatbot = gr.Chatbot(label="Awesome airlines chatbot assistant")
    msg = gr.Textbox(label="message (API is turned off, message me to turn it back on)", placeholder="")
    clear = gr.ClearButton([msg, chatbot])

    def respond(message, chat_history: list):
        bot_message = query(message, chat_history)
        chat_history.append((message, bot_message))
        return "", chat_history

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

demo.launch()