import aiohttp
import gradio as  gr
import numba
import requests 
import  base64
from PIL import Image
import  io
import json
from numba import jit
import matplotlib.pyplot as plt


examples =  ["examples/0002_01_00_01_55.jpg",
"examples/0-spoof.jpg",
"examples/0.jpg",
"examples/3.jpg",
"examples/6-mask.jpg",
"examples/AGL752VM_id147_s0_150.png",
"examples/FT720P_G780_REDMI4X_id0_s0_105.png",
"examples/7.jpg"]


async def spoof_trigger(b64):
    url = "https://spoofapi1.azurewebsites.net/api/spoofvisualize"
    payload = {"img": b64}
    headers = {
    'x-functions-key': 'wGw3zXXPlLCez-VrcSs9RTahE4gLC674pf7Fp6Au2kUHAzFuNnZZMw==',
    'Content-Type': 'text/plain'
    }

    async with aiohttp.ClientSession() as session:
        async with session.post(url, json=payload, headers=headers) as response:
            response_text = await response.text()
    return response_text
# @jit
async def predict_image(img):
    # Convert NumPy array to PIL Image
    img = Image.fromarray(img.astype('uint8'))
    
    # Create a BytesIO object
    buffer = io.BytesIO()
    
    # Save the PIL Image to the BytesIO object
    img.save(buffer, format='JPEG')
    
    # Get the base64 representation
    img_base64 = base64.b64encode(buffer.getvalue()).decode()

    print(len(img_base64))
    # # img_base64 to  plot
    # img = Image.open(io.BytesIO(base64.b64decode(img_base64)))
    # # img save 
    # img.save("img.jpg")


    res = await spoof_trigger(img_base64)
    # print(json.loads(res))
    spoof_res = json.loads(res)['spoof_res']
    annotated_image  = json.loads(res)['annotated_image']
    conf_score = float( json.loads(spoof_res)['confidence_score'])

    # img_base64 to  plot
    img = Image.open(io.BytesIO(base64.b64decode(annotated_image)))
    # img save 
    img.save("cache/img.jpg")

    confidences = {'Real': 1-conf_score, 'Fake': conf_score}

    return (confidences,annotated_image)


with gr.Blocks(title="Spoof-Demo", css="#custom_header {min-height: 3rem; text-align:  center} #custom_title {min-height: 3rem; text-align: center}") as demo :
    gr.Markdown("## Face Antispoof-Demo", elem_id="custom_title")
    gr.Markdown("Gradio Demo for Face Antispoofing", elem_id="custom_header")
    gr.Markdown("👨‍💻 Only fot research preview Intended")
    with gr.Row():
        with gr.Column():
            with gr.Box():
                gr.Markdown("### Input")
                image = gr.Image(label="Input Image")
                image.style(height=240) 
                btn = gr.Button(text="Submit")
                btn.style(full_width=True)          
        with gr.Column():
            with gr.Box():
                gr.Markdown("### Output")
                output_image = gr.Image(label="Output Image")
                output_image.style(height=240) 
                label_probs  =  gr.outputs.Label()

    btn.click(predict_image, image , outputs=[label_probs,output_image ],api_name="Face Antispoofing")
    gr.Examples(
        examples=examples,
        inputs=image,
        outputs = output_image,   
        fn=predict_image,
        cache_examples=False,
    ) 

if  __name__ == "__main__":
    demo.launch(debug=True)