import gradio as gr
import os
import requests
import json

api_key = os.environ.get("OPENAI_API_KEY")

def openai_chat(prompt, history):

    if not prompt:
        return gr.update(value='', visible=len(history) < 10), [(history[i]['content'], history[i+1]['content']) for i in range(0, len(history)-1, 2)], history

    url = "https://api.openai.com/v1/chat/completions"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }
    prompt_msg = {
                "role": "user",
                "content": prompt
            }
    
    data = {
        "model": "gpt-3.5-turbo",
        "messages": history + [prompt_msg]
    }

    response = requests.post(url, headers=headers, json=data)
    json_data = json.loads(response.text)

    response = json_data["choices"][0]["message"]

    history.append(prompt_msg)
    history.append(response)

    return gr.update(value='', visible=len(history) < 10), [(history[i]['content'], history[i+1]['content']) for i in range(0, len(history)-1, 2)], history

def start_new_conversation(history):
    history = []
    return gr.update(value=None, visible=True), gr.update(value=[]), history

css = """
      #col-container {max-width: 50%; margin-left: auto; margin-right: auto;}
      #chatbox {min-height: 400px;}
      #header {text-align: center;}
      """

with gr.Blocks(css=css) as demo:
    history = gr.State([])

    with gr.Column(elem_id="col-container"):
        gr.Markdown("""## OpenAI Chatbot Demo
                    Using the ofiicial API and GPT-3.5 Turbo model.
                    Current limit is 5 messages per conversation.""",
                    elem_id="header")
        chatbot = gr.Chatbot(elem_id="chatbox")
        input_message = gr.Textbox(show_label=False, placeholder="Enter text and press enter", visible=True).style(container=False)
        btn_start_new_conversation = gr.Button("Start a new conversation")

    input_message.submit(openai_chat, [input_message, history], [input_message, chatbot, history])
    btn_start_new_conversation.click(start_new_conversation, [history], [input_message, chatbot, history])
            
if __name__ == "__main__":
    demo.launch(debug=True, height='800px')