import os
import gradio as gr
from io import BytesIO
import zipfile
import tempfile
import random  
import json
from gradio_client import Client, handle_file
import base64
from PIL import Image

BUY_PREMIUM = "🥇 Get Premium Token to Unlock"

STATUS_MESSAGES = {
    201: "Not satisfied with the results? Get a Premium Token for instant, deeper results including social profiles.", 
    301: "Invalid token! Get Premium Token using link in page"
}

backend = Client(os.getenv("BACKEND"))

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

def search_face(file, token=None):
    result_text = backend.predict(
            file=handle_file(file),
            token=token, 
            api_name="/search_face"
    )
   
    result = json.loads(result_text)
    outarray = []
    if result['status'] > 300:
        raise gr.Error(STATUS_MESSAGES[result['status']])
    
    for item in result['result']:
        image = Image.open(BytesIO(base64_to_image(item['image'])))
        outarray.append((image, item['url']))
    
    if result['status'] == 201:
        gr.Info(STATUS_MESSAGES[result['status']], duration=11)
    return outarray

def export_images(items):
    if not items:
        return None
    # Create a zip file in memory
    zip_buffer = BytesIO()
    with zipfile.ZipFile(zip_buffer, 'w') as zip_file:
        url_text = ""
        i = 1
        for item in items:
            if item[1] == BUY_PREMIUM:
                continue
            with open(item[0], 'rb') as img_file:
                zip_file.writestr(f'image_{i}.jpg', img_file.read())
            url_text += f"image_{i}.jpg: {item[1]}\n"
            i += 1
        zip_file.writestr("urls.txt", url_text)
    zip_buffer.seek(0)

    with tempfile.NamedTemporaryFile(delete=False, suffix=".zip") as temp_file:
        temp_file.write(zip_buffer.getvalue())
        temp_file_path = temp_file.name

    return temp_file_path

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

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

div#component-9 {
    background-color: tomato;
}

.svelte-snayfm {
    height: auto;
}

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

js = """
function aff() {
const links = document.querySelectorAll('a');

const currentUrl = new URL(window.location.href);
const currentParams = currentUrl.searchParams.toString();

links.forEach(link => {
    const href = link.getAttribute('href');
    if (href && (href.startsWith('https://faceonlive.pocketsflow.com') || href.startsWith('https://faceonlive.com'))) {
        const targetUrl = new URL(href);
        // Append current page parameters to the link
        currentParams.split('&').forEach(param => {
            if (param) {
                const [key, value] = param.split('=');
                targetUrl.searchParams.set(key, value);
            }
        });
        link.setAttribute('href', targetUrl.toString());
    }
});

  return ''
}
"""

head = """
<!-- Google tag (gtag.js) -->
<script async src="https://www.googletagmanager.com/gtag/js?id=G-8YPXF4536P"></script>
<script>
  window.dataLayer = window.dataLayer || [];
  function gtag(){dataLayer.push(arguments);}
  gtag('js', new Date());

  gtag('config', 'G-8YPXF4536P');
</script>
"""

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)

def search_face_examples(image):
    return search_face(image), gr.update(visible=False), gr.update(visible=True)

MARKDOWN0 = """
    # Free Face Search Online - ❤️Like above if this space helps
    #### [Learn more about our Reverse Face Search](https://faceonlive.com/face-search-online)
    #### [Check out our Face Search API integration and Affiliate Program here.](https://portfolio.faceonlive.com/#face_search)
"""
MARKDOWN1 = """
    ### [Get Premium Token - Deep Search & Social Media Access](https://faceonlive.pocketsflow.com/checkout?productId=676c15b1971244a587ca07cb)  
    ### [Protect Your Privacy – Start Your Takedown Now](https://faceonlive.pocketsflow.com/checkout?productId=676d7e7597f8b3b699f84820)
"""
MARKDOWN2 = """
### [Why Get the Premium Token?](https://faceonlive.pocketsflow.com/checkout?productId=676c15b1971244a587ca07cb)  
✅ **Deep Search**  
✅ **Social Media and Deep Web Scan**  
✅ **Exclusive Access to Hidden Profiles**  
✅ **No Subscription Needed**
"""
MARKDOWN3 = """
<div align="right"><a href="https://faceonlive.com/reverse-image-search" target='_blank' style='font-size: 16px;'>Looking For Reverse Image Search</div>
"""
with gr.Blocks(css=custom_css, head=head, delete_cache=(3600, 3600)) as demo:
    gr.Markdown(MARKDOWN0)
    with gr.Row():
        with gr.Column(scale=1) as col1:
            image = gr.Image(type='filepath', height=360)
            with gr.Row():
              with gr.Column():
                token = gr.Textbox(placeholder="(Optional) Get Premium Token Below.", label="Premium Token")
                gr.Markdown(MARKDOWN1)
              with gr.Column():
                md_premium1 = gr.Markdown(MARKDOWN2)
            with gr.Row():
                search_face_button = gr.Button("Reverse Face Search")
            gr.HTML(MARKDOWN3)
            gr.Examples(['examples/1.jpg', 'examples/2.jpg'], inputs=image, cache_examples=True, fn=search_face_examples, outputs=[output, col1, col2])
            
        with col2.render():
            gr.Markdown("> ### **⚠️ Reminder:** Export images before refreshing the page by clicking the 'Export Images' button.")
            output.render()
            with gr.Row():
                export_button = gr.Button("Export Images")
                export_file = gr.File(label="Download")
            with gr.Row():
                with gr.Column():
                    gr.Markdown(MARKDOWN1)
                with gr.Column():
                    md_premium2 = gr.Markdown(MARKDOWN2)
            new_search_button = gr.Button("🚀 Try New Search")
            gr.HTML(MARKDOWN3)

    search_face_button.click(search_ui, inputs=[], outputs=[col1, col2], api_name=False)
    search_face_button.click(search_face, inputs=[image, token], outputs=[output], api_name=False)
    export_button.click(export_images, inputs=[output], outputs=export_file, 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%2FFaceOnLive%2FFace-Search-Online"><img src="https://api.visitorbadge.io/api/visitors?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2FFaceOnLive%2FFace-Search-Online&labelColor=%23ff8a65&countColor=%2337d67a&style=flat&labelStyle=upper" /></a>')
    html = gr.HTML()
    demo.load(None, inputs=None, outputs=html, js=js)

demo.queue(api_open=False, default_concurrency_limit=8).launch(show_api=False)