import gradio as gr
import openai
import os

current_dir = os.path.dirname(os.path.abspath(__file__))
css_file = os.path.join(current_dir, "style.css")

initial_prompt = "You are a helpful assistant."


def parse_text(text):
    lines = text.split("\n")
    for i,line in enumerate(lines):
        if "```" in line:
            items = line.split('`')
            if items[-1]:
                lines[i] = f'<pre><code class="{items[-1]}">'
            else:
                lines[i] = f'</code></pre>'
        else:
            if i>0:
                line = line.replace("<", "&lt;")
                line = line.replace(">", "&gt;")
                lines[i] = '<br/>'+line.replace(" ", "&nbsp;")
    return "".join(lines)


def get_response(system, context, raw=False):
    openai.api_key = "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"
    response = openai.Completion.create(
        engine="text-davinci-002",
        prompt=f"{system}{''.join([f'{c['role']}: {c['content']}\n' for c in context])}",
        max_tokens=1024,
        n=1,
        stop=None,
        temperature=0.5,
    )

    message = response.choices[0].text
    message_with_stats = f"{message}"
    return message, parse_text(message_with_stats)


def predict(input_sentence):
    if len(input_sentence) == 0:
        return []

    chatbot.append((input_sentence, message_with_stats))

    context.append({"role": "user", "content": f"{input_sentence}"})

    message, message_with_stats = get_response(systemPrompt.value["content"], context)

    context.append({"role": "assistant", "content": message})

    return chatbot, context


def retry():
    if len(context) == 0:
        return [], []

    context[-1]["content"] = "Could you rephrase that?"

    message, message_with_stats = get_response(systemPrompt.value["content"], context[:-1])

    context[-1] = {"role": "assistant", "content": message}

    chatbot[-1] = (context[-2]["content"], message_with_stats)

    return chatbot, context


def delete_last_conversation():
    if len(context) == 0:
        return [], []

    chatbot = chatbot[:-1]
    context = context[:-2]
    return chatbot, context


def reduce_token():
    context.append({"role": "user", "content": "Please summarize our conversation and reduce tokens used. Don't include this prompt."})

    response = get_response(systemPrompt.value["content"], context, raw=True)

    optmz_str = f'Okay, we talked about: {response.choices[0].text}\n\nTotal tokens used this conversation: {response.choices[0].logprobs.top_logprobs[0].tokens}'

    chatbot.append(("Please summarize our conversation and reduce tokens used. Don't include this prompt.", parse_text(optmz_str)))
    context = [{"role": "assistant", "content": f"Okay, we talked about: {response.choices[0].text}"}]
    return chatbot, context


def reset_state():
    return [], []


def update_system(new_system_prompt):
    return {"role": "system", "content": new_system_prompt}


title = """<h1 align="center">You Ask, I Answer - Chatbot</h1>"""

description = "This chatbot is designed to assist you with any questions or tasks you may have. Simply type in your query and the chatbot will provide you with a helpful response."

systemPrompt = gr.inputs.Textbox(lines=2, label="Enter the system prompt you would like to use:")
userInput = gr.inputs.Textbox(lines=2, label="Enter your message:")

chatbot_output = gr.outputs.HTML(type="auto")

chatbot_interface = gr.Interface(
predict,
[systemPrompt, userInput],
chatbot_output,
title=title,
description=description,
theme="compact",
layout="vertical",
examples=[
["Can you help me with my math homework?", "Sure, what do you need help with?"],
["How can I make pizza from scratch?", "First, you will need to gather the ingredients..."]
],
article="https://openai.com/blog/how-to-build-a-state-of-the-art-conversational-ai-with-transfer-learning-2021/"
)

if name == "main":
chatbot_interface.launch(debug=True)