from dotenv import load_dotenv
load_dotenv("docker/.env")

from pymongo.mongo_client import MongoClient
from pymongo.server_api import ServerApi
from src.mongo_search import MongoSearch
from src.image_search import ImageSearcher

import os
import gradio as gr

uri = os.getenv("MONGO_URI")
db = os.getenv("MONGO_DB")
cars_collection = os.getenv("MONGO_CARS_COLLECTION")
image_collection = os.getenv("MONGO_IMAGE_COLLECTION")
img_search_index = os.getenv("MONGO_SEARCH_IMG_INDEX")
car_search_index = os.getenv("MONGO_SEARCH_CAR_INDEX")

client = MongoClient(uri, server_api=ServerApi('1'))
db = client[db]
cars_collection = db[cars_collection]
image_collection = db[image_collection]

mongo_search = MongoSearch(cars_collection, car_search_index, index_variable="review_embedding") #desc_embedding
image_searcher = ImageSearcher(image_collection, 
                                cars_collection, 
                                img_search_index)

def search_mongo(search_query):
    object = mongo_search(search_query)
    return  object[0]['makeModel'], object[0]['photoUrls'], \
            object[1]['makeModel'], object[1]['photoUrls'], \
            object[2]['makeModel'], object[2]['photoUrls']


def search_img(image):

    result = image_searcher(image)['results']
    return result[0]['makeModel'], result[0]['photoUrls']



with gr.Blocks() as demo:

    with gr.Tab("Text search"):
        with gr.Column():
            with gr.Row():
                with gr.Column():
                    text_output_img_1 = gr.Gallery(height=300, selected_index=0)
                    text_output_text_1 = gr.Textbox(label='Output car model', show_label=True)

                with gr.Column():
                    text_output_img_2 = gr.Gallery(height=300, selected_index=0)
                    text_output_text_2 = gr.Textbox(label='Output car model', show_label=True)

                with gr.Column():
                    text_output_img_3 = gr.Gallery(height=300, selected_index=0)
                    text_output_text_3 = gr.Textbox(label='Output car model', show_label=True)

        text_input_text = gr.Textbox(label='Search Query', show_label=True)
        text_btn  = gr.Button()


    with gr.Tab("Image search"):
        with gr.Row():
            with gr.Column():
                input_img = gr.Image(type='pil')
                img_btn  = gr.Button()
            with gr.Column():
                img_output_img = gr.Gallery(preview=True)
                img_output_text = gr.Textbox(label='Output car model', show_label=True)


    text_btn.click(search_mongo,
                    inputs=text_input_text,
                    outputs=[text_output_text_1, text_output_img_1,
                             text_output_text_2, text_output_img_2,
                             text_output_text_3, text_output_img_3])
    
    

    img_btn.click(search_img,
                inputs=input_img,
                outputs=[img_output_text, img_output_img])


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