import PIL.Image
import gradio as gr
import huggingface_hub
import onnxruntime as rt
import numpy as np
import cv2
from PIL import ImageOps
import requests
from pathlib import Path
from PIL import Image
import base64
import numpy as np
import io
import os
# import streamlit as st

print(os.environ.get("ENDPOINT"))


def get_mask(img_in):

    print(f"{img_in=}")

    # covert to 8-bit RGB
    img_in = img_in.convert("RGB")

    img_in_smol = img_in.resize((256, 256))

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

    upload_url = st.secrets["ENDPOINT"]

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

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

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

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

        return (img_in, result['emotion'])
    else:
        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)