import gradio as gr
import time
from utils import Bot
from utils.functions import make_documents, make_descriptions

def init_bot(file=None,title=None,pdf=None,key=None):
    
    if key is None:
        return 'You must submit OpenAI key'
    
    if pdf is None:
        return 'You must submit pdf file'
    
    if file is None:
        return 'You must submit media file'
    
    if title is None:
        return 'You must submit the description of the media'
    
    file = file.name
    print(file)
    pdf = pdf.name
    file_description = make_descriptions(file, title)
    # print(file_description)
    documents = make_documents(pdf)
    
    # print(documents[0])
    global bot
    
    bot = Bot(
        openai_api_key=key,
        file_descriptions=file_description,
        text_documents=documents,
        verbose=False
    )
    
    return 'Chat bot successfully initialized'

def msg_bot(history):
    message = history[-1][0]
    
    bot_message = bot(message)['output']
    history[-1][1] = ""
    for character in bot_message:
        history[-1][1] += character
        time.sleep(0.05)
        yield history

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

with gr.Blocks() as demo:
    
    key = gr.Textbox(label='OpenAI key')
    with gr.Tab("Chat bot initialization"):
            
        with gr.Row(variant='panel'):
            with gr.Column():
                with gr.Row():
                    title = gr.Textbox(label='File short description')
                with gr.Row():
                    file = gr.File(label='CSV or image', file_types=['.csv', 'image'])

            pdf = gr.File(label='pdf')
            
        with gr.Row(variant='panel'):
            init_button = gr.Button('submit')
            init_output = gr.Textbox(label="Initialization status")
            init_button.click(fn=init_bot,inputs=[file,title,pdf,key],outputs=init_output,api_name='init')
            
        chatbot = gr.Chatbot()
        msg = gr.Textbox(label='Ask the bot')
        clear = gr.Button('Clear')
        msg.submit(user,[msg,chatbot],[msg,chatbot],queue=False).then(
            msg_bot, chatbot, chatbot
        )
        clear.click(lambda: None, None, chatbot, queue=False)

demo.queue()
demo.launch()