from utils.utils import get_logger, initialization, get_search_result, get_image_description
import gradio as gr
import logging


logger = get_logger()
collection = None


def main(query):
    logger = logging.getLogger(__name__)
    print("Starting search...")
    logger.info("Starting search...")
    print("-------------------------------------------------------")
    logger.info("-------------------------------------------------------")
    exit = False
    while not exit:
        if query == "exit":
            exit = True
            print("-------------------------------------------------------")
            logger.info("-------------------------------------------------------")
            print("Search terminated.")
            logger.info("Search terminated.")
            return None, "Search terminated."
        else:
            # Get search result including the original description of the image
            image, text = get_search_result(collection, data_set, query, model, n_results=2)

            # Get detailed description of the image generated by multimodal LLM
            description_by_llm = get_image_description(image)
            
            return image, text, description_by_llm


if __name__ == "__main__":
    try:
        if collection == None:
            collection, data_set, model, logger = initialization(logger)
        
        app = gr.Interface(
            fn=main,
            inputs=[gr.Textbox(label="Describe the scene that you are looking for:")],
            outputs=[gr.Image(label="Here's the scene found based on your input:"),
                     gr.Textbox(label="Original description of the found scene:"),
                     gr.Textbox(label="Detailed description generated by LLM:")],
            title="Search for a scene in the world of GTA!"
        )
        app.launch(share=True)
    except Exception as e:
        logger.exception(e)
        raise e