import gradio as gr
import io
import requests
import os
import base64
import cv2
from PIL import Image
import json
import numpy as np
import rsa

def FaceSwap(source_img, dest_img):

    secret = os.environ['API_KEY']
    privateKey = os.environ["PRIVATE_KEY"]
    publicKey = os.environ["PUBLIC_KEY"]
    baseurl = os.environ["BASE_URL"]

    n,e,d,p,q = privateKey.split(",")
    n1,e1 = publicKey.split(",")
    privateKey = rsa.PrivateKey(int(n),int(e),int(d),int(p),int(q))
    publicKey = rsa.PublicKey(int(n1), int(e1))

    encMessage = rsa.encrypt(secret.encode('utf-8'), publicKey)
    encMessage = base64.b64encode(encMessage).decode("utf-8")
    
    src_size = source_img.size
    dest_size = dest_img.size
    
    src_base64 = base64.b64encode(bytearray(source_img.tobytes())).decode('utf-8')
    dest_base64 = base64.b64encode(bytearray(dest_img.tobytes())).decode('utf-8')

    #Sending request to a FaceSwap API I deployed it on AWS ec2. 
    #Contact maytamaku.saidhwar@gmail.com to get access to the API.
    api_url = f"{baseurl}/api/faceswap"
    
    payload = {"source_img": str(src_base64), "dest_img": str(dest_base64), "source_size": src_size, "dest_size": dest_size}
    headers = {'Content-Type': 'multipart/form-data', "Authorization": encMessage}
    res = requests.post(api_url, json=payload, headers=headers)

    if (res.status_code >= 200 and res.status_code <= 299):
        res = np.array(res.json())
        return res
    elif (res.status_code >= 400 and res.status_code <= 499):
        raise gr.Error(f"Code {res.status_code}: Erron on client")
    else:
        raise gr.Error(f"Code {res.status_code}: Erron on server. Try after some time")
    
if __name__ == "__main__":
    with gr.Blocks(title="Face Swap | Meet With Your new Personality") as demo:
        gr.Interface(
            fn=FaceSwap,
            inputs=[
                gr.Image(type="pil"),
                gr.Image(type="pil")
            ],
            outputs=[
                gr.Image(
                    type="pil",
                    show_download_button=False,
                    show_share_button=False)
            ],
            allow_flagging="never"
            )
    demo.queue().launch(show_api=False, favicon_path="./favicon.ico")