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: 4")

    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, prevText):
        bot_message = generateAIMessage(history, prevText, 0)
        history[-1][1] = bot_message
        return history

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

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

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

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

    def generateAIMessage(history, prevText, characterId):
        if checkIfQuestionsLeft(prevText):
            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
        else:
            return "--- YOU'VE RUN OUT OF QUESTIONS, MAKE YOUR CHOICE ---"

    def updateQuestions(prevText):
        remainingQuestions = int(prevText["label"].split(": ")[1])
        remainingQuestions -= 1
        if remainingQuestions > 0:
            return "Remaining Questions: " + str(remainingQuestions)
        else:
            return "You've used up your questions, now choose the culprit!"
    
    def checkIfQuestionsLeft(prevText):
        remainingQuestions = int(prevText["label"].split(": ")[1])
        return remainingQuestions > 0

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

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

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

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

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


iface.launch()