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

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 base64.b64decode(base64_str + '=' * (-len(base64_str) % 4))

def check_db(img_array, suffix, token_server_url):
    hashes = []
    out_array = []
    for item in img_array:
        try:
            image_data = base64_to_image(item["image"])
            
            hash_object = hashlib.sha256(image_data)
            image_hash = hash_object.hexdigest()
            
            hashes.append(image_hash)
            out_array.append((Image.open(BytesIO(image_data)), item["url"] + suffix))
    
        except base64.binascii.Error as e:
            raise ValueError(f"Invalid base64 string: {str(e)}")
        except Exception as e:
            raise ValueError(f"Error processing image: {str(e)}")

    try:
        r = requests.post(url=token_server_url + '/lookup-hashes', json={"hashes": hashes})
        out_array = [value for i, value in enumerate(out_array) if i not in r.json().get('res')]
    except:
        raise gr.Error("Token Server Error!")

    return out_array

def search_face(file, token):
    free_url = os.environ.get("SERVER_URL_FREE")
    premium_url = os.environ.get("SERVER_URL_PREMIUM")
    token_server_url = os.environ.get("TOKEN_URL")

    url = free_url
    if token:
        try:
            r = requests.post(url=token_server_url + '/verify-token', json={"token": token})
            if r.json().get('status') == 'success':
                url = premium_url
            else:
                raise gr.Error("Invalid token! For free search, use empty string for token")
        except:
            raise gr.Error("Invalid token!")

    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("Please try again later.")
    elif status_code == 404:
        gr.Info("Timeout, try again.")

    if status_code > 300:
        return []   

    try:        
        res = r.json().get('img_array')
        suffix = "*********"
        if url == premium_url:
            suffix = ""
        
        out_array = check_db(res, suffix, token_server_url)

        if url == premium_url:
            try:
                r = requests.post(url=token_server_url + '/activate-token', json={"token": token})
            except:
                raise gr.Error("Invalid token!")

        return out_array
    except:
        raise gr.Error("Try again.")

def search_face_free(file):
    return search_face(file, '')

custom_css = """
caption.caption {
    user-select: text;
    cursor: text;
}
"""
with gr.Blocks(css=custom_css) 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>
    """
    )
    with gr.Row():
        with gr.Column(scale=1):
            image = gr.Image(type='filepath', height=480)
            token = gr.Textbox(placeholder="(Optional) Premium Token via the link below.", label="Premium Token")
            gr.HTML("<a href='https://faceonlive.pocketsflow.com/checkout?productId=676c15b1971244a587ca07cb' target='_blank'>Get Premium Token: Perform Deep Search & Full URLs</a>")            
            gr.HTML("<a href='https://faceonlive.pocketsflow.com/checkout?productId=676d7e7597f8b3b699f84820' target='_blank'>Send DMCA & GDPR Takedown Notices On Your Behalf, Opt-Out From Search</a>")
            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")
    
    gr.Examples(['examples/1.jpg', 'examples/2.jpg'], inputs=image, cache_examples=True, cache_mode='lazy', fn=search_face_free, outputs=[output])

    search_face_button.click(search_face, inputs=[image, token], outputs=[output], 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)