import gradio as gr
import openai
import os
import requests

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
ADMIN_PASSWORD = os.getenv("ADMIN_PASSWORD")

openai.api_key = OPENAI_API_KEY

default_system_message = {"role": "system", "content": "You are a brilliant, helpful assistant, always providing answers to the best of your knowledge. If you are unsure of the answer, you indicate it to the user. Currently, you don't have access to the internet."}

def get_completion(model, user_message, message_history, chatlog_history, temperature, maximum_length, top_p, frequency_penalty, presence_penalty):
    new_history_row = {"role": "user", "content": user_message}
    updated_message_history = message_history + [new_history_row]
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {openai.api_key}",
    }
    payload = {
        "model":model,
        "messages":updated_message_history,
        "temperature":temperature,
        "max_tokens":maximum_length,
        "top_p":top_p,
        "frequency_penalty":frequency_penalty,
        "presence_penalty":presence_penalty,
    }
    completion = requests.post(
        "https://api.openai.com/v1/chat/completions",
        headers=headers,
        json=payload,
    )
    completion = completion.json()
    # completion = openai.ChatCompletion.create(
    #     model=model,
    #     messages=updated_message_history,
    #     temperature=temperature,
    #     max_tokens=maximum_length,
    #     top_p=top_p,
    #     frequency_penalty=frequency_penalty,
    #     presence_penalty=presence_penalty,
    # )
    assistant_message = completion["choices"][0]["message"]["content"]
    new_history_row = {"role": "assistant", "content": assistant_message}
    updated_message_history = updated_message_history + [new_history_row]
    updated_chatlog_history = chatlog_history + [(user_message, assistant_message)]
    token_count = completion["usage"]["total_tokens"]
    return "", updated_message_history, updated_chatlog_history, updated_chatlog_history, token_count

def retry_completion(model, message_history, chatlog_history, temperature, maximum_length, top_p, frequency_penalty, presence_penalty):
    # get latest user message
    user_message = chatlog_history[-1][0]
    # delete latest entries from chatlog history
    updated_chatlog_history = chatlog_history[:-1]
    # delete latest assistant message from message_history
    updated_message_history = message_history[:-1]
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {openai.api_key}",
    }
    payload = {
        "model":model,
        "messages":updated_message_history,
        "temperature":temperature,
        "max_tokens":maximum_length,
        "top_p":top_p,
        "frequency_penalty":frequency_penalty,
        "presence_penalty":presence_penalty,
    }
    completion = requests.post(
        "https://api.openai.com/v1/chat/completions",
        headers=headers,
        json=payload,
    )
    completion = completion.json()
    # completion = openai.ChatCompletion.create(
    #     model=model,
    #     messages=updated_message_history,
    #     temperature=temperature,
    #     max_tokens=maximum_length,
    #     top_p=top_p,
    #     frequency_penalty=frequency_penalty,
    #     presence_penalty=presence_penalty,
    # )
    assistant_message = completion["choices"][0]["message"]["content"]
    new_history_row = {"role": "assistant", "content": assistant_message}
    updated_message_history = updated_message_history + [new_history_row]
    updated_chatlog_history = updated_chatlog_history + [(user_message, assistant_message)]
    token_count = completion["usage"]["total_tokens"]
    return "", updated_message_history, updated_chatlog_history, updated_chatlog_history, token_count

def reset_chat():
    return "", [default_system_message], [], [], 0

theme = gr.themes.Monochrome()
with gr.Blocks(theme=theme) as app:
    message_history = gr.State([default_system_message])
    chatlog_history = gr.State([])
    with gr.Row():
        with gr.Column(scale=4):
            chatbot = gr.Chatbot(label="Chat").style(height=600)
        with gr.Column(scale=1):
            # model = gr.Textbox(lines=3, value="You are a brilliant, helpful assistant, always providing answers to the best of your knowledge. If you are unsure of the answer, you indicate it to the user.", interactive=True, label="System")
            model = gr.Dropdown(choices=["gpt-3.5-turbo", "gpt-4", "gpt-4-32k"], value="gpt-3.5-turbo", interactive=True, label="Model")
            temperature = gr.Slider(minimum=0, maximum=1, step=0.05, value=0.5, interactive=True, label="Temperature")
            maximum_length = gr.Slider(minimum=0, maximum=2048, step=64, value=256, interactive=True, label="Maximum length")
            top_p = gr.Slider(minimum=0, maximum=1, step=0.01, value=1, interactive=True, label="Top P")
            frequency_penalty = gr.Slider(minimum=0, maximum=2, step=0.01, value=0, interactive=True, label="Frequency penalty")
            presence_penalty = gr.Slider(minimum=0, maximum=2, step=0.01, value=0, interactive=True, label="Presence penalty")
            token_count = gr.Number(interactive=False, label="Token Count")
    with gr.Row():
        user_message = gr.Textbox(label="Message")
    with gr.Row():
        reset_button = gr.Button("Reset Chat")
        retry_button = gr.Button("Retry")

    user_message.submit(get_completion, inputs=[model, user_message, message_history, chatlog_history, temperature, maximum_length, top_p, frequency_penalty, presence_penalty], outputs=[user_message, message_history, chatlog_history, chatbot, token_count])
    retry_button.click(retry_completion, inputs=[model, message_history, chatlog_history, temperature, maximum_length, top_p, frequency_penalty, presence_penalty], outputs=[user_message, message_history, chatlog_history, chatbot, token_count])
    reset_button.click(reset_chat, inputs=[], outputs=[user_message, message_history, chatlog_history, chatbot, token_count])

app.launch(auth=("admin", ADMIN_PASSWORD))