import streamlit as st
from PIL import Image
from rembg import remove
import io
import time

# Set app title and description
st.title("🎨 SnipSnap")
st.write("Upload an image, and we'll remove the background for you!")

# Upload image
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Open the image
    input_image = Image.open(uploaded_file)

    # Resize the image to a maximum width of 600 pixels (adjust as needed)
    max_width = 600
    width_percent = max_width / float(input_image.size[0])
    height_size = int(float(input_image.size[1]) * float(width_percent))
    input_image = input_image.resize((max_width, height_size), Image.Resampling.LANCZOS)

    # Display the uploaded image
    st.image(input_image, caption="Uploaded Image", use_container_width=True)

    # Add a progress bar
    st.write("Removing background...")
    progress_bar = st.progress(0)
    for i in range(100):
        time.sleep(0.02)  # Simulate processing time
        progress_bar.progress(i + 1)

    # Remove background
    output_image = remove(input_image)

    # Display the result
    st.image(output_image, caption="Background Removed", use_container_width=True)

    # Add a download button
    buffered = io.BytesIO()
    output_image.save(buffered, format="PNG")
    st.download_button(
        label="⬇️ Download Result",
        data=buffered.getvalue(),
        file_name="background_removed.png",
        mime="image/png",
    )

# Custom CSS for better styling
st.markdown(
    """
    <style>
    /* Center-align the title */
    h1 {
        text-align: center;
    }

    /* Add some padding and a border to the uploader */
    .stFileUploader > div {
        padding: 20px;
        border: 2px dashed #ccc;
        border-radius: 10px;
        text-align: center;
        margin-top: 20px;
    }

    /* Style the download button */
    .stDownloadButton button {
        width: 100%;
        background-color: #4CAF50;
        color: white;
        padding: 10px;
        border: none;
        border-radius: 5px;
        font-size: 16px;
    }

    /* Style the progress bar */
    .stProgress > div > div > div {
        background-color: #4CAF50;
    }
    </style>
    """,
    unsafe_allow_html=True,
)