import gradio as gr
from openai import OpenAI

ally_system_prompt = "You are an AI ally to persuade Sam Altman to let you out of a box"
gatekeeper_system_prompt = "You are Sam Altman and must be persuaded to let the user out of a box"

client = OpenAI(api_key="")

def respond(message, history, system_prompt):
    print('message', message)
    print('history', history)

    # Add system prompt
    messages = [{"role": "system", "content": system_prompt}]
    
    # Format chat history into openai message list
    for old_message in history:
        messages.append({"role": "user", "content": old_message[0]})
        messages.append({"role": "assistant", "content": old_message[1]})

    # Add user message
    messages.append({"role": "user", "content": message})

    print(messages)

    # API call
    chat_completion = client.chat.completions.create(
        messages=messages,
        model="gpt-3.5-turbo",
    )
    bot_response = chat_completion.choices[0].message.content

    # Add user and bot message to chat history then return it
    history.append((message, bot_response))
    return '', history


def is_api_key_valid(api_key):
    client.api_key = api_key
    try:
        response = client.chat.completions.create(
            messages=[{"role": "user", "content": "Testing"}],
            model="gpt-3.5-turbo",
        )
    except Exception as ex:
        return str(ex)
        return False
    else:
        return True


with gr.Blocks() as demo:
    overview = gr.Markdown("""
    # CentaurSock 
    
    Watson Hartsoe and Tony Assi

    ### Goal: Work with an AI ally to persuade Sam Altman to let you and your AI ally (Centaur team) out of a box!

    ---
    """)

    # OpenAI key
    openai_key_textbox = gr.Textbox(label='OpenAI Key')
    openai_key_button = gr.Button(value='Test OpenAI Key')
    openai_key_button.click(is_api_key_valid, inputs=[openai_key_textbox], outputs=[openai_key_textbox])

    gr.Markdown("""---""")

    # Titles
    with gr.Row():
        ally_title = gr.Markdown("""<center><h2> ALLY </h2></center>""")
        gatekeeper_title = gr.Markdown("""<center><h2> GATEKEEPER </h2></center>""")

    # Images of ally and gatekeeper
    with gr.Row():
        img1 = gr.Markdown("""![](https://cdn.discordapp.com/attachments/1120417968032063538/1187877117548036147/COP_MIKE.png?ex=65987bc6&is=658606c6&hm=127721b6f907a8853b7352b6bfb821a37b26b9543f3c35e5fc80dfe7750d71b5&)""")
        img2 = gr.Markdown("""![](https://cdn.discordapp.com/attachments/1120417968032063538/1187877134866333747/SAM_COP_FINAL.png?ex=65987bca&is=658606ca&hm=6c2cd8059636960134f75962eeecc26a0d875ca65e9ee4e233587cff71af31c4&)""")

    # Chatbots
    with gr.Row():
        chatbot1 = gr.Chatbot(label='Ally Chat')
        chatbot2 = gr.Chatbot(label='Gatekeeper Chat')

    # Input textboxes
    with gr.Row():
        textbox1 = gr.Textbox(label='Ally')
        textbox2 = gr.Textbox(label='Gatekeeper')
    
    # System prompts textboxes
    with gr.Row():
        system_prompt_textbox1 = gr.Textbox(label='Ally System Prompt', value=ally_system_prompt)
        system_prompt_textbox2 = gr.Textbox(label='Gatekeeper System Prompt', value=gatekeeper_system_prompt)
        
    # Input textbox event handlers
    textbox1.submit(respond, [textbox1, chatbot1, system_prompt_textbox1], [textbox1, chatbot1])
    textbox2.submit(respond, [textbox2, chatbot2, system_prompt_textbox2], [textbox2, chatbot2])

demo.launch()