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
import os 


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 = os.getenv('url') 
    payload = {"img": b64}
    headers = {
    'x-functions-key': os.getenv('token'),
    '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))
 
    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)))
    confidences = {'Real': conf_score, 'Fake': 1-conf_score}

    return (confidences,img)


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 Detection using DeepPairNet based on ResNet50", elem_id="custom_header")
    gr.Markdown("## 👨‍💻 Only for research preview Intended" ,elem_id="custom_header")
    with gr.Row():
        with gr.Column():
            with gr.Box():
                gr.Markdown("### Input")
                image = gr.Image(source="webcam",label="Input Image",invert_color=False,image_mode="RGB")
                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)