import openai
import gradio as gr
import os
import logging
import json

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

openai.api_key = os.environ['key']
initMsg = os.environ['init']
roleAns = os.environ['roleAns']
classAsk = os.environ['classAsk']
classAns = os.environ['classAns']
defaultMsg = os.environ['defaultMsg']

mod = 3;

def trimMessages(messages):
    # Assuming messages is an array of objects with "role" and "content" properties
    content = {}
    for message in messages:        
        if message["role"] == "assistant":
            if "Current Status:" in message["content"]:
                content = message["content"]
                start_index = message["content"].index("Current Status:") + len("Current Status:")
                end_index = message["content"].index("Wielding:")
                new_content = message["content"][:start_index] + message["content"][end_index:]
                message["content"] = new_content
    if content != {}:
        messages.append({"role": "assistant", "content": content})
    return messages

def getChooseRole(msg):
    start_index = msg.index("As an ") + len("As an ")
    end_index = msg.index(",")
    return msg[start_index:] + msg[:start_index]

def getChooseClass(msg):
    full = getChooseRole(msg)
    start_index = full.index(" ") + len(" ")
    end_index = msg.index(",")
    return full[start_index:] + full[:start_index]

def chatbot(input, messages):
    messages = messages or [{"role": "system", "content": initMsg}]
    if len(messages) == 1:
        
        memory = trimMessages(messages)
        logging.info("put memory:"+printMessages(memory))
        chat = openai.ChatCompletion.create(
            model="gpt-3.5-turbo", messages=memory,
            max_tokens=1048,n=1,temperature=0.5,
        )
        logging.info("put memory return:"+chat.choices[0].message.content)

    if input:
        messages.append({"role": "user", "content": input})
        chat = openai.ChatCompletion.create(
            model="gpt-3.5-turbo", messages=messages[-2:],
            max_tokens=2048,n=1,temperature=0.5,
        )
        reply = chat.choices[0].message.content

        ans = ""
        if len(messages) == 3:
            ans = "I choose "+ getChooseRole(reply) +" as my race." + classAsk
            logging.info("ChooseRole:"+ans)
        if len(messages) == 5:
            ans = "My Race is ，I choose "+ getChooseClass(reply) +" as my class." + classAns
            logging.info("ChooseClass:"+ans)
        messages.append({"role": "assistant", "content": ans})
        return reply, messages

def printMessages(messages):
    delimiter = '\n'
    msg_string = delimiter.join([f"{obj['role']}:{obj['content']}" for obj in messages])
    logging.info("messages:"+msg_string)
    return msg_string

app = gr.Interface(fn=chatbot, inputs=[gr.Textbox(lines=7, label="You ask and answer questions below"), "state"], 
                   outputs=[gr.Textbox(label="DND Game Reply", placeholder=roleAns), "state"], title="DND Game",#, gr.Textbox(label="History"), 
                   description="DND Game",theme="compact")
app.launch(share=False)