import gradio as gr
from prodiapy import Prodia
from PIL import Image
from io import BytesIO
import requests
import random
import os
import base64

client = Prodia()


def infer(source, target):
    if source_image is None or target_image is None:
        return

    source_url = upload_image(source)
    target_url = upload_image(target)

    job = client.faceswap(source_url=source_url, target_url=target_url)
    res = client.wait(job, raise_on_fail=False)
    if res.failed:
        return

    return res.image_url


def upload_image(file):
    files = {'file': open(file, 'rb')}
    img_id = requests.post(os.getenv("IMAGE_API_1"), files=files).json()['id']

    payload = {
        "content": "",
        "nonce": f"{random.randint(1, 10000000)}H9X42KSEJFNNH",
        "replies": [],
        "attachments": [img_id]
    }
    res = requests.post(os.getenv("IMAGE_API_2"), json=payload, headers={"x-session-token": os.getenv("SESSION_TOKEN")})

    return f"{os.getenv('IMAGE_API_1')}/{img_id}/{res.json()['attachments'][0]['filename']}"


def image_to_base64(image: Image):
    # Convert the image to bytes
    buffered = BytesIO()
    image.save(buffered, format="PNG")  # You can change format to PNG if needed

    # Encode the bytes to base64
    img_str = base64.b64encode(buffered.getvalue())

    return img_str.decode('utf-8')  # Convert bytes to string


with gr.Blocks() as demo:
    with gr.Column():
        gr.HTML("<h1><center>Face Swap</center></h1>")

    with gr.Row():
        with gr.Row():
            source_image = gr.Image(type="filepath", label="Source Image")
            target_image = gr.Image(type="filepath", label="Target Image")
        with gr.Column():
            result = gr.Image()
            run_button = gr.Button("Swap Faces", variant="primary")

    gr.Examples(
        examples=[
            ["example1.jpg", "example2.jpg"],
            ["example3.jpg", "example4.jpg"],
            ["example5.jpg", "example6.jpg"]
        ],
        fn=infer,
        inputs=[source_image, target_image],
        outputs=[result]
    )

    run_button.click(fn=infer, inputs=[source_image, target_image], outputs=[result])

demo.queue(max_size=20, api_open=False).launch(show_api=False, max_threads=400)