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
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))

    # 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)

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

    background_img = background_img.convert("RGBA")

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

    input_img = input_img.convert("RGBA")

    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

# 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:        
        
        background_img_file = background_img_file.read()
        
        cropped_pic = st_cropperjs(pic=background_img_file, btn_text="Generate!", 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
            img = process_image(cropped_pic)

            # Display the combined image
            st.image(img, caption="Combined Image", 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",
                data=img_bytes,
                file_name="free_palestine.png",
                mime="image/png",
            )

if __name__ == "__main__":
    main()