import streamlit as st
from PIL import Image, ImageOps
import io

from rembg import remove
import requests
from PIL import Image
from io import BytesIO
import os

from streamlit_cropperjs import st_cropperjs

# Function to process the image

# 1st picture: normal behaviour

# 2nd picture: normal picture and applying the palestinian flag with an opacity of 120

def process_image(image):

    print(image)

    os.makedirs('original', exist_ok=True)
    os.makedirs('masked', exist_ok=True)
    image = image.convert("RGB")

    image.save('original/image.jpg', format = "jpeg")

    output_path = "masked/image.png"

    with open(output_path, "wb") as f:
        input = open('original/image.jpg', 'rb').read()
        subject = remove(input)
        f.write(subject)

    palestine_bg = "https://flagdownload.com/wp-content/uploads/Flag_of_Palestine_Flat_Round-1024x1024.png"
    background_img = Image.open(BytesIO(requests.get(palestine_bg).content))

    background_img = background_img.resize((image.width, image.height))
    #image = image.resize((background_img.width, background_img.height))

    image = image.convert("RGBA")

    background_img = background_img.convert("RGBA")

    input_img = Image.open('original/image.jpg')

    input_img = input_img.convert("RGBA")

    normal_img = Image.blend(image, background_img, .3).convert('RGB')

    # Create a semi-transparent overlay
    overlay = Image.new('RGBA', background_img.size, (0, 0, 0, 120))  # Adjust the last value (120) to control transparency
    background_img = Image.alpha_composite(background_img.convert('RGBA'), overlay)

    combined_img = Image.alpha_composite(input_img, background_img)

    combined_img = combined_img.convert('RGB')
    # combined_img.save('masked/finale.jpg', format='jpeg')

    foreground_img = Image.open(output_path)
    combined_img.paste(foreground_img, (0,0), foreground_img)

    combined_img = combined_img.convert('RGB')
    # combined_img.save("masked/background_final.jpg", format="jpeg")

    return combined_img, normal_img

# Streamlit app
def main():
    st.title("Watermelon PFP Generator - by Kaito")

    # Select background image
    background_img_file = st.file_uploader("Select your image", type=["jpg", "png"])
    if background_img_file is not None: 

        # handle error if not picture
        if background_img_file.type.split("/")[0] != "image":
            st.error("Please upload an image file")
            return       
        
        background_img_file = background_img_file.read()
        
        cropped_pic = st_cropperjs(pic=background_img_file, btn_text="Generate!", size=1.0, key="foo")
        if cropped_pic:
            # save the cropped pic
            cropped_pic = Image.open(io.BytesIO(cropped_pic))
            cropped_pic.save("cropped_pic.png")

            # Process the images

            with st.spinner('Generating your fabulous 🍉 profile picture...'):
                img, img2 = process_image(cropped_pic)

            # Create two columns
            col1, col2 = st.columns(2)

            # Display the images in each column
            with col1:
                st.image(img, caption="Image 1", use_column_width=True)

                img_bytes = BytesIO()
                img.save(img_bytes, format='PNG')
                img_bytes = img_bytes.getvalue()

                # Add a download button
                st.download_button(
                    label="Download Image 1",
                    data=img_bytes,
                    file_name="free_palestine.png",
                    mime="image/png",
                )

            with col2:
                st.image(img2, caption="Image 2", use_column_width=True)
                                
                img_bytes = BytesIO()
                img2.save(img_bytes, format='PNG')
                img_bytes = img_bytes.getvalue()

                # Add a download button
                st.download_button(
                    label="Download Image 2",
                    data=img_bytes,
                    file_name="free_palestine2.png",
                    mime="image/png",
                )



if __name__ == "__main__":
    main()