import streamlit as st
import requests
import os
from PIL import Image
import shutil
import zipfile
from io import BytesIO

# Streamlit app title
st.title("Background Remover App")

# Input for API Key
api_key = st.text_input("Enter Stability API Key", type="password")

# File uploader for images
uploaded_files = st.file_uploader("Upload Images", type=["png", "jpg", "jpeg", "webp"], accept_multiple_files=True)

# Output format selector
output_format = st.selectbox("Select Output Format", options=["png", "webp"], index=0)

# Process button
if st.button("Start Processing"):
    if not api_key:
        st.error("Please enter your API key.")
    elif not uploaded_files:
        st.error("Please upload at least one image.")
    else:
        progress_bar = st.progress(0)
        processed_files = []
        
        # Create a temporary directory to store processed files
        temp_dir = "processed_images"
        os.makedirs(temp_dir, exist_ok=True)
        
        for i, uploaded_file in enumerate(uploaded_files):
            image = uploaded_file.read()
            try:
                response = requests.post(
                    "https://api.stability.ai/v2beta/stable-image/edit/remove-background",
                    headers={
                        "authorization": f"Bearer {api_key}",
                        "accept": "image/*"
                    },
                    files={"image": image},
                    data={"output_format": output_format},
                )
                
                if response.status_code == 200:
                    output_file_path = os.path.join(temp_dir, f"{uploaded_file.name.split('.')[0]}_processed.{output_format}")
                    with open(output_file_path, "wb") as f:
                        f.write(response.content)
                    
                    processed_files.append(output_file_path)
                    st.success(f"Processed: {uploaded_file.name}")
                    st.image(Image.open(output_file_path), caption=f"Processed: {uploaded_file.name}")
                else:
                    st.error(f"Failed to process {uploaded_file.name}: {response.json()}")
            except Exception as e:
                st.error(f"Error processing {uploaded_file.name}: {e}")
            progress_bar.progress((i + 1) / len(uploaded_files))
        
        if processed_files:
            # Create a zip archive of all processed images
            zip_buffer = BytesIO()
            with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
                for file_path in processed_files:
                    zip_file.write(file_path, os.path.basename(file_path))
            zip_buffer.seek(0)

            # Provide download link for the zip file
            st.success("Processing completed!")
            st.download_button(
                label="Download All Processed Images as ZIP",
                data=zip_buffer,
                file_name="processed_images.zip",
                mime="application/zip"
            )
        
        # Cleanup temporary directory
        shutil.rmtree(temp_dir)