from omegaconf import OmegaConf
from query import VectaraQuery
import os
import gradio as gr


def isTrue(x) -> bool:
    if isinstance(x, bool):
        return x
    return x.strip().lower() == 'true'

corpus_ids = str(os.environ['corpus_ids']).split(',')
cfg = OmegaConf.create({
    'customer_id': str(os.environ['customer_id']),
    'corpus_ids': corpus_ids,
    'api_key': str(os.environ['api_key']),
    'title': os.environ['title'],
    'description': os.environ['description'],
    'source_data_desc': os.environ['source_data_desc'],
    'streaming': isTrue(os.environ.get('streaming', False)),
    'prompt_name': os.environ.get('prompt_name', None)
})


vq = VectaraQuery(cfg.api_key, cfg.customer_id, cfg.corpus_ids, cfg.prompt_name)


def respond(message, history):
    if cfg.streaming:
        # Call stream response and stream output
        stream = vq.submit_query_streaming(message)

        outputs = ""
        for output in stream:
            outputs += output
            yield outputs
    else:
        # Call non-stream response and return message output
        response = vq.submit_query(message)
        yield response

cfg.title = f'''<h1 style="text-align: center;">{cfg.title}</h1>'''

cfg.description = f'''<center> <h2>{cfg.description}</h2>
                        <br>
                        This demo uses Retrieval Augmented Generation to ask questions about {cfg.source_data_desc}</center>'''

image_html = '''
<img src="https://github.com/david-oplatka/chatbot-streamlit/blob/main/Vectara-logo.png" style="float:left; width:100px; height:50px; margin-right:20px;">
'''

# with gr.Blocks() as demo:
#     gr.HTML()
#     gr.ChatInterface(respond, title = cfg.title, description = cfg.description,
#                         chatbot = gr.Chatbot(value = [[None, "How may I help you?"]], scale=3))

with gr.Blocks() as demo:
    with gr.Row():
        gr.HTML(image_html)
    gr.HTML(cfg.title)
    gr.HTML(cfg.description)
    gr.ChatInterface(respond, chatbot=gr.Chatbot(value=[[None, "How may I help you?"]], scale=3))


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