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="JPEG", quality=90)
    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, headers={"X-RapidAPI-Key": os.environ.get("API_KEY")}, json={"image": image_base64})
    except:
        raise gr.Error("Please select image file!")

    status_code = r.status_code
    if status_code == 301:
        gr.Info("Too many faces in the photo.")
    elif status_code == 302:
        gr.Info("Face is not clear enough.")
    elif status_code == 303:
        gr.Info("No matches found.")
    elif status_code == 304:
        gr.Info("No face in the photo.")
    elif status_code == 305:
        gr.Info("Search blocked for privacy issue.")
    elif status_code == 401:
        gr.Info("Invalid image format.")
    elif status_code == 402:
        gr.Info("Wrong request.")
    elif status_code == 403:
        gr.Info("Requests all used in your token.")
    elif status_code == 404:
        gr.Info("Timeout, try again.")

    if status_code > 300:
        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/>
    ## [FaceOnLive: On-premises ID Verification, Biometric Authentication Solution Provider](https://faceonlive.com)    
    <br>
    
    ## Looking to embed this on your website? Use the HTML code below.
    ```html
    <script type="module" src="https://gradio.s3-us-west-2.amazonaws.com/5.5.0/gradio.js"></script>
    <gradio-app src="https://faceonlive-face-search-online.hf.space"></gradio-app>
    ```
    """
    )
    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, default_concurrency_limit=4).launch(server_name="0.0.0.0", server_port=7860, show_api=False)