import streamlit as st
import os
from datetime import datetime
from PIL import Image
from io import BytesIO

from src.utils import change_background, matte
from src.st_style import apply_prod_style

# apply_prod_style(st)  # NOTE: Uncomment this for production!


def image_download_button(pil_image, filename: str, fmt: str, label="Download"):
    if fmt not in ["jpg", "png"]:
        raise Exception(f"Unknown image format (Available: {fmt} - case sensitive)")
    
    pil_format = "JPEG" if fmt == "jpg" else "PNG"
    file_format = "jpg" if fmt == "jpg" else "png"
    mime = "image/jpeg" if fmt == "jpg" else "image/png"
    
    buf = BytesIO()
    pil_image.save(buf, format=pil_format)
    
    return st.download_button(
        label=label,
        data=buf.getvalue(),
        file_name=f'{filename}.{file_format}',
        mime=mime,
    )


st.title("AI Photo Background Removal")
st.image(Image.open("assets/demo.jpg"))
st.write(
    """
    You want to remove your photo background, but don't have the time and effort to learn photo editing skills?
    **This app will change or remove your photo background, in seconds.**
    """
)

uploaded_file = st.file_uploader(
    label="Upload your photo here", 
    accept_multiple_files=False, type=["png", "jpg", "jpeg"],
)

if uploaded_file is not None:

    with st.expander("Original photo", expanded=True):
        if uploaded_file is not None:
            st.image(uploaded_file)
        else:
            st.warning("You haven't uploaded any photo yet")

    in_mode = st.selectbox("Choose background color", ["Transparent (PNG)", "White", "Black", "Green", "Red", "Blue"])
    in_submit = st.button("Submit")

    if uploaded_file is not None and in_submit:
        img_input = Image.open(uploaded_file)
        
        with st.spinner("AI is doing magic to your photo. Please wait..."):
            hexmap = {
                "Transparent (PNG)": "#000000",
                "Black": "#000000",
                "White": "#FFFFFF",
                "Green": "#22EE22",
                "Red": "#EE2222",
                "Blue": "#2222EE",
            }
            alpha = 0.0 if in_mode == "Transparent (PNG)" else 1.0
            img_matte = matte(img_input)
            img_output = change_background(img_input, img_matte, background_alpha=alpha, background_hex=hexmap[in_mode])
        
        with st.expander("Success!", expanded=True):
            st.image(img_output)
            uploaded_name = os.path.splitext(uploaded_file.name)[0]
            image_download_button(
                pil_image=img_output,
                filename=uploaded_name,
                fmt="png"
            )