import numpy as np
import gradio as gr
from z_generate import ServerlessInference
from z_embedding import load_vector_store

# STATIC TEXT DISPLAY
TXT_APP_DESCRIPTION = '''
Just another Retrieval Augmented Generation that also retrieves images
'''
TXT_SOURCE_DOC_DESCRIPTION = '''
Manually Downloaded as HTML files: 

1. https://en.wikipedia.org/wiki/MS_Dhoni
2. https://en.wikipedia.org/wiki/Jharkhand
2. https://en.wikipedia.org/wiki/Cricket_World_Cup

## Details

1. Vector Store is built using FAISS prior to starting this app. Although the vector store size in KBs but the creation and loading of the store takes processing takes ~10GB RAM and lasts 5 mins. Hence **NOT BUILDING IT DURING RUNTIME OF APP**.

'''


# UI Interface
demo = gr.Blocks()
vector_text, vector_image = load_vector_store()
llm = ServerlessInference(vector_store_text=vector_text, vector_store_images=vector_image)

# Processing Functions
def update_response(query:str = "something"):
    response_text = llm.perform_rag(query)
    return response_text
    
def update_gallery(text:str = "hell"):
    imgs = [
                ("http://www.marketingtool.online/en/face-generator/img/faces/avatar-1151ce9f4b2043de0d2e3b7826127998.jpg", "Some Description"),
                ("http://www.marketingtool.online/en/face-generator/img/faces/avatar-116b5e92936b766b7fdfc242649337f7.jpg", "Another Description")
            ]
    return imgs


def ask_bot(text):
    return update_response(text), update_gallery(text)

# UI Layout
with demo:
    gr.Markdown(TXT_APP_DESCRIPTION)

    with gr.Tabs():
        with gr.TabItem("Ask Bot"):
            with gr.Row(equal_height=True):
                with gr.Column(scale=3):
                    text_input = gr.Textbox(
                        label="You query here",
                        placeholder="What positions apart from crickter did Dhoni held?"
                    )
                with gr.Column(scale=1):
                    btn_generate = gr.Button("Generate Answer")

            with gr.Row():
                with gr.Column(scale=3):
                    text_output = gr.Textbox(label="Bot Response:", placeholder="Type in Query before I could answer")
                
                with gr.Column(scale=2): 
                    gallery = gr.Gallery(
                        label="Generated images", show_label=False, elem_id="gallery"
                        , columns=[3], rows=[1], object_fit="contain", height="auto"
                    )

            btn_generate.click(ask_bot, text_input, outputs=[text_output, gallery])



        ####
        with gr.TabItem("Source Documents"):
            gr.Markdown(TXT_SOURCE_DOC_DESCRIPTION)

demo.launch(debug=True)