import gradio as gr
import openai

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}\n\n{context}",
        temperature=0.7,
        max_tokens=1024,
        top_p=1,
        frequency_penalty=0,
        presence_penalty=0,
        stop=None,
    )
    if raw:
        return response
    else:
        statistics = f'This conversation Tokens usage【{response["total_characters"]} / 2048】 （ Question + above {response["prompt"]["length"]}，Answer {response["choices"][0]["length"]} ）'
        message = response.choices[0].text

        message_with_stats = f'{message}\n\n{statistics}'

        return message, parse_text(message_with_stats)

def predict(chatbot, input_sentence, system, context):
    if len(input_sentence) == 0:
        return []
    context.append(input_sentence)

    message, message_with_stats = get_response(system, context)

    context.append(message)

    chatbot.append((input_sentence, message_with_stats))

    return chatbot, context

def retry(chatbot, system, context):
    if len(context) == 0:
        return [], []
    message, message_with_stats = get_response(system, context[:-1])
    context[-1] = message

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

def delete_last_conversation(chatbot, context):
    if len(context) == 0:
        return [], []
    chatbot = chatbot[:-1]
    context = context[:-2]
    return chatbot, context

def reset_state():
    return [], []

def update_system(new_system_prompt):
    return new_system_prompt

title = """<h1 align="center">Tu întrebi și eu răspund.</h1>"""
description = """<div align=center>

Will not describe your needs to ChatGPT？You Use [ChatGPT Shortcut](https://newzone.top/chatgpt/)

</div>
"""

gr.Interface(
    predict, 
    inputs=["textbox", "text", "text", "text"],
    outputs="text",
    title=title,
    description=description,
    allow_flagging=False,
    theme="huggingface",
    layout="vertical",
    examples=[["What's your name?", "I don't have a name. How can I assist you today?", "", ""]],
    interpretation="default"
).launch()