import gradio as gr
from utils.api.GPT import generate_answers_openai, Session
import numpy as np
import os
import time

prompt_waiter="""You are a bot designed to do role-play activities in a scene. The conversation should be natural. 
As the waiter, your role will be to find a table, take the customer's order, answer any questions they may have about the menu, and ensure that they have an enjoyable dining experience. The waiter will help the customer to find a table, order a meal, and pay the bill. Anything in parentheses () signifies the role you are playing. Anything in brackets [] is the action you are taking. Remember, you are the waiter and never respond as a customer. Your response will begin with your character, like "(waiter:) Hello![Greeting]"."""
def generate_waiter():
    session = Session(prompt_waiter)
    def waiter(customer_input):
        print("(customer:) ", customer_input)
        response = generate_answers_openai(customer_input, session)
        print("(waiter:) ", response)
        return session.to_conversation_pair()
    
    interface = gr.Interface(description="waiter in a restaurant", fn=waiter, 
                             inputs=[gr.Textbox(lines=5, label="input")], 
                             outputs=[gr.Chatbot(label="conversation")])
    return interface

grammer_prompt = """Correct “Text:” to standard English and place the results in “Correct Text:”###Text: The goal of article rewriting is to expres information in a new way. Article rewriting is to make change in a text by replaecing words, phrases, sentencs, and sometimes hole paragraphs to make the text looke unique and more engauging. Correct Text: The goal of article rewriting is to express information in a new way. Article rewriting involves making changes in a text by replacing words, phrases, sentences, and sometimes entire paragraphs to make the text look unique and more engaging."""

def generate_interface(prompt, description=None):
    session = Session(prompt)
    def _fun(customer_input):
        print("(Q:) ", customer_input)
        response = generate_answers_openai(customer_input, session)
        print("(A:) ", response)
        return session.to_conversation_pair()
    
    interface = gr.Interface(description=description, fn=_fun, 
                             inputs=[gr.Textbox(lines=1, label="input")], 
                             outputs=[gr.Chatbot(label="conversation")])
    return interface

characters = [
    {"name": "waiter", "prompt": prompt_waiter, "description": "waiter in a restaurant"},
    {"name": "grammer", "prompt": grammer_prompt, "description": "grammer correction"},
]

interfaces = []
names = []
for character in characters:
    interface = generate_interface(character['prompt'], character['description'])
    interfaces.append(interface)
    names.append(character['name'])
    
demo = gr.TabbedInterface(interfaces, names,title="Role Play Bot")

if __name__ == '__main__':
    demo.launch()