import PIL.Image
import gradio as gr
from PIL import ImageOps
import requests
from pathlib import Path
from PIL import Image
import base64
import numpy as np
import io
import os
from random import choice
from dotenv import load_dotenv

load_dotenv()


def get_mask(img_in):

    # print(f"{img_in=}")
    img_in = img_in.convert("RGB")

    file_path = "/tmp/img_in.jpg"
    img_in.save(file_path)

    upload_url = os.environ.get("ENDPOINT")

    files = {'file': open(file_path, 'rb')}

    response = requests.post(upload_url, files=files)

    if response.status_code == 200:

        if 'error' in response.json():
            print(f"ERROR: {response.json()['error']}")
            gr.Error(response.json()['error'])
            return (None, None)

        result = response.json()
        print('Result:', result)

        all_bgs = list(Path("examples").glob("*.jpg"))

        bg_img = Image.open(choice(all_bgs)).convert(
            "RGBA").resize(img_in.size)

        bg_img = ImageOps.fit(
            bg_img,
            img_in.size,
            Image.ANTIALIAS
        )

        mask = Image.open(io.BytesIO(base64.b64decode(
            result['mask']))).resize(img_in.size)

        img_in = ImageOps.autocontrast(img_in, cutoff=0.1).convert("RGBA")
        img_in.putalpha(mask)

        img_in = Image.alpha_composite(bg_img, img_in)

        return (img_in, result['emotion'])
    else:
        gr.Error(response.text)
        # print('error:', response.text)


footer = r"""
<center>
<b>
DEMO FOR BG REMOVAL
</b>
</center>
"""

with gr.Blocks(title="Face Shine") as app:
    gr.HTML("<center><h1>ARKA Remove Background</h1></center>")
    with gr.Row():
        with gr.Column():
            input_img = gr.Image(type="pil", label="Input image")
            # new_img = gr.Image(type="pil", label="Custom background")
            run_btn = gr.Button(variant="primary")
        with gr.Column():
            output_img = gr.Image(type="pil", label="result")
            emotion = gr.Label(label="emotion")

    run_btn.click(get_mask, [input_img], [output_img, emotion])

    with gr.Row():
        gr.HTML(footer)

app.launch(share=False, debug=True, enable_queue=True, show_error=True)