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')}
    headers = {"Authorization": f"Bearer {os.environ.get('PERMITTED_TOKEN')}"}

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

    if response.status_code == 200:
        result = response.json()

        img_square = Image.open(io.BytesIO(base64.b64decode(
            result['square'])))
        # img_stories = Image.open(io.BytesIO(base64.b64decode(
        #     result['stories'])))

        return (img_square, result['emotion'])
    else:
        print(f"ERROR: {response.json()['detail']}")
        raise gr.Error(response.json()['detail'])


header = r"""
<center>
<h1>ARKA Selfie App</h1>
<p>Send this app a selfie to add a custom background and elements to it.</p>
<p>It also detects and displays human emotions and works anywhere.</p>
</center>
"""

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

with gr.Blocks(title="ARKA + TikTok Demo") as app:
    gr.HTML(header)
    with gr.Row():
        run_btn = gr.Button(variant="primary", interactive=True, label="Run")
    with gr.Row():
        with gr.Column():
            input_img = gr.Image(type="pil", label="Input image")
        with gr.Column():
            with gr.Row():
                output_img = gr.Image(type="pil", label="result_square")
            emotion = gr.Label(label="emotion")

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

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

app.launch(share=True, debug=True, enable_queue=True, show_error=True,
           auth=((os.environ.get("USERNAME"), os.environ.get("PASSWORD"))))