import os
import gradio as gr
from PIL import Image
from gradio_client import Client, handle_file
import json

BACKEND = os.getenv('BACKEND')
TOKEN = os.getenv('TOKEN')
backend = Client(BACKEND, hf_token=TOKEN)

def search_image(file):    
    result_text = backend.predict(
            file=handle_file(file),
            api_name="/search_image"
    )
   
    result = json.loads(result_text)
    outarray = []
    for item in result['result']:
        outarray.append((item['image'], item['url']))
    return outarray

custom_css = """
caption.caption {
    user-select: text;
    cursor: text;
}

div#component-18 {
    max-height: 63.39px;
}

.svelte-snayfm {
    height: auto;
}

.icon.svelte-snayfm {
    width: 48px;
    height: 48px;
}
"""

output = gr.Gallery(label="Found Images (The search may take a couple of minutes.)", columns=[3], object_fit="contain", height="480px", interactive=False)
col2 = gr.Column(scale=2, visible=False)

def init_ui():
    return gr.update(visible=True), gr.update(visible=False)

def search_ui():
    return gr.update(visible=False), gr.update(visible=True)

with gr.Blocks(css=custom_css, delete_cache=(3600, 3600)) as demo:
    gr.Markdown(
        """
    # Free Reverse Image Search
    ### [Learn more about our Reverse Image Search](https://faceonlive.com/reverse-image-search)
    <br/>
    """
    )
    with gr.Row():
        with gr.Column(scale=1) as col1:
            image = gr.Image(type='filepath', height=360)          
            gr.HTML("<a href='https://faceonlive.com/face-search-online' target='_blank' style='font-size: 20px;'>Try Reverse Face Search Here</a>")  
            search_image_button = gr.Button("Reverse Image Search")            
            
        with col2.render():
            output.render()
            gr.HTML("<a href='https://faceonlive.com/face-search-online' target='_blank' style='font-size: 20px;'>Try Reverse Face Search Here</a>")
            new_search_button = gr.Button("🚀 Try New Search")
        
    search_image_button.click(search_ui, inputs=[], outputs=[col1, col2], api_name=False)
    search_image_button.click(search_image, inputs=[image], outputs=[output], api_name=False)

    new_search_button.click(init_ui, inputs=[], outputs=[col1, col2], api_name=False)

    gr.HTML('<a href="https://visitorbadge.io/status?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2FReverseImageSearch%2FReverse-Image-Search"><img src="https://api.visitorbadge.io/api/visitors?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2FReverseImageSearch%2FReverse-Image-Search&labelColor=%23ff8a65&countColor=%2337d67a&style=flat&labelStyle=upper" /></a>')

demo.queue(api_open=False, default_concurrency_limit=8).launch(server_name="0.0.0.0", server_port=7860, show_api=False)