import os
import gradio as gr
import requests
import base64
from io import BytesIO
from PIL import Image

count = 0
def image_to_base64(image):
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode('utf-8')

def base64_to_image(base64_str):
    return Image.open(BytesIO(base64.b64decode(base64_str + '=' * (-len(base64_str) % 4))))

def search_face(file):
    global count    
    url = os.environ.get("SERVER_URL")

    try:
        image = Image.open(file)
        image_base64 = image_to_base64(image)
        r = requests.post(url=url, json={"token": os.environ.get("ACCESS_TOKEN"), "type": "novip", "image": image_base64})
    except:
        raise gr.Error("Please select image file!")
    
    if r.text == "No matches":
        gr.Info("No images found.")
        return [], count
    
    try:
        res = r.json().get('img_array')
        out_array = []
        for item in res:
            out_array.append((base64_to_image(item["image"]), item["url"] + "*********"))
        count += 1            
        return out_array, count
    except:
        raise gr.Error("Try again.")


with gr.Blocks() as demo:
    gr.Markdown(
        """
    # Search Your Face Online For Free
    ## For more detailed information, please check on our website.<br/>https://faceonlive.com<br/><br/>
    ## For premium support or partnership inquiries, contact us.<br/><br/>
    """
    )
    with gr.Row():
        with gr.Column(scale=1):
            image = gr.Image(type='filepath', height=480)
            search_face_button = gr.Button("Search Face")
        with gr.Column(scale=2):
            output = gr.Gallery(label="Found Images", columns=[4], object_fit="contain", height="auto")
            countwg = gr.Number(label="Count")
    
    gr.Examples(['examples/1.jpg', 'examples/2.jpg'], inputs=image, cache_examples=True, cache_mode='lazy', fn=search_face, outputs=[output, countwg])

    search_face_button.click(search_face, inputs=image, outputs=[output, countwg], api_name=False)

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


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