import gradio as gr
import requests
from PIL import Image
import io
import base64
import uuid

def image_base64(image):
    
    if image is None:
        return None

    img_byte_arr = io.BytesIO()
    image.save(img_byte_arr, format='JPEG')
    img_byte_arr = img_byte_arr.getvalue()

    base64_image = base64.b64encode(img_byte_arr).decode('utf-8')

    return base64_image

def process_images(Image, reference_image):

   #create two unique image names using uuid
    ref_img = uuid.uuid4().hex
    user_img  = uuid.uuid4().hex


    model_image = [{'name': f"ref_img_{ref_img}.JPEG", 'image': image_base64(reference_image)}]
    user_image = [{'name': f"user_img_{user_img}.JPEG", 'image': image_base64(Image)}]
    
    req = {
    "input": {
        "model_workflow": "default_wf",
        "result_encoding_format": "base64",
        "model_encoding_format": "base64",
        "user_encoding_format": "base64",
        "model_image": model_image,
        "user_photo": user_image
        }
    }

    URI = "https://api.fotolabs.app/runsync"



    response = requests.post(URI, json=req, headers = {
        "Authorization": f"Bearer {'Y9IXJ0O3QBIY1T6M49GG9LHDF83F8Y51BTV2978C'}"
    })

    if response.status_code == 200:
        data = response.json()
        task_id = data.get('id')
        img_str = response.json()['output']['message'][0]
        output_image = base64_to_img(img_str)
        print(task_id)
        return output_image
    else:
        return None

def img_to_base64(img):
    buffer = io.BytesIO()
    img.save(buffer, format="JPEG")
    img_str = base64.b64encode(buffer.getvalue()).decode("utf-8")
    return img_str

def base64_to_img(img_str):
    img_bytes = base64.b64decode(img_str)
    img = Image.open(io.BytesIO(img_bytes))
    return img

# Create Gradio interface
iface = gr.Interface(
    fn=process_images,
    inputs=[gr.Image(label="Image", type="pil"),
            gr.Image(label="Reference Image", type="pil")],
    outputs=gr.Image(label="Output Image"),
    
).queue(default_concurrency_limit=5)
iface.queue()
# Launch the Gradio interface
iface.launch()
# iface.launch()