import gradio as gr
import os
import openai

openai.api_key = "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"

with gr.Blocks() as iface:
    chatbot = []
    msg = []
    roles = ["Security Guard", "Curator", "Researcher", "Conservationist", "Guide"]
    initText = ["You are Steve Binner, 31, a security guard for a museum of medieval "
                "history. You're absolutely sure that this place is haunted. A month ago, "
                "before the spooky stuff started, you were really happy with your job, "
                "but that changed real fast. The ghost is a french nobleman named Baron "
                "Hugo Antoine the Third. You've been trying to tell other people about "
                "this, but  you've only been laughed at so far. But you've got evidence! "
                "Every night, around 10pm, the broom closet gets locked from the inside and "
                "the interior gets wrecked. We're talking pushed over tables, "
                "broken containers, the whole shebang. When it comes to the murder of "
                "director Eisenholz last thursday, that has to have been the ghost, "
                "no question. You even have the door card records for the room he was in: "
                "After the museum closed at 5pm, nobody entered that room till the next "
                "morning. Now you're being interrogated by a detective. You don't use "
                "uptight language, and you're not super well educated on most stuff, "
                "but when it comes to the paranormal, you're an ace. ",
                "You are a Curator. ",
                "You are a Researcher. ",
                "You are a Conservationist. ",
                "You are a Guide. "]

    i = 0
    while i < len(roles):
        with gr.Tab(roles[i]):
            chatbot.append(gr.Chatbot())
            msg.append(gr.Textbox())
        i += 1

    questionCounter = gr.Label(value="Remaining Questions: 24")

    def user(user_message, history):
        return "", history + [[user_message, None]]

    # this is an ugly workaround, you cannot pass regular integers to gradio functions, so i had to do it this way
    def bot0(history):
        bot_message = generateAIMessage(history, 0)
        history[-1][1] = bot_message
        return history

    def bot1(history):
        bot_message = generateAIMessage(history, 1)
        history[-1][1] = bot_message
        return history

    def bot2(history):
        bot_message = generateAIMessage(history, 2)
        history[-1][1] = bot_message
        return history

    def bot3(history):
        bot_message = generateAIMessage(history, 3)
        history[-1][1] = bot_message
        return history

    def bot4(history):
        bot_message = generateAIMessage(history, 4)
        history[-1][1] = bot_message
        return history

    def generateAIMessage(history, characterId):
        message_history = [
            {"role": "system", "content": initText[characterId] + "Stay in character. Use natural language. Don't "
                                                                  "reveal all of the information in a single message,"
                                                                  " and leave hints. "}
        ]
        for pair in history:
            message_history.append({"role": "user", "content": pair[0]})
            print()
            if pair[1] is not None:
                message_history.append({"role": "assistant", "content": pair[1]})
        completion = openai.ChatCompletion.create(
            model="gpt-3.5-turbo",
            messages=message_history
        )
        return completion.choices[0].message.content

    def updateQuestions(prevText):
        remainingQuestions = int(prevText["label"].split(": ")[1])
        remainingQuestions -= 1
        return "Remaining Questions: " + str(remainingQuestions)

    # again, ugly workaround
    msg[0].submit(user, [msg[0], chatbot[0]], [msg[0], chatbot[0]], queue=False).then(
        bot0, chatbot[0], chatbot[0]
    ).then(
        updateQuestions, questionCounter, questionCounter
    )

    msg[1].submit(user, [msg[1], chatbot[1]], [msg[1], chatbot[1]], queue=False).then(
        bot1, chatbot[1], chatbot[1]
    ).then(
        updateQuestions, questionCounter, questionCounter
    )

    msg[2].submit(user, [msg[2], chatbot[2]], [msg[2], chatbot[2]], queue=False).then(
        bot2, chatbot[2], chatbot[2]
    ).then(
        updateQuestions, questionCounter, questionCounter
    )

    msg[3].submit(user, [msg[3], chatbot[3]], [msg[3], chatbot[3]], queue=False).then(
        bot3, chatbot[3], chatbot[3]
    ).then(
        updateQuestions, questionCounter, questionCounter
    )

    msg[4].submit(user, [msg[4], chatbot[4]], [msg[4], chatbot[4]], queue=False).then(
        bot4, chatbot[4], chatbot[4]
    ).then(
        updateQuestions, questionCounter, questionCounter
    )


iface.launch()