import os
import zipfile
import shutil
from PIL import Image
import io
from rembg import remove
import gradio as gr
from concurrent.futures import ThreadPoolExecutor

def resize_and_crop_image(image_path, target_size=(1080, 1080), crop_mode='center'):
    with Image.open(image_path) as img:
        width, height = img.size

        # Calculate the scaling factor
        scaling_factor = max(target_size[0] / width, target_size[1] / height)

        # Resize the image with high-quality resampling
        new_size = (int(width * scaling_factor), int(height * scaling_factor))
        resized_img = img.resize(new_size, Image.LANCZOS)

        if crop_mode == 'center':
            left = (resized_img.width - target_size[0]) / 2
            top = (resized_img.height - target_size[1]) / 2
        elif crop_mode == 'top':
            left = (resized_img.width - target_size[0]) / 2
            top = 0
        elif crop_mode == 'bottom':
            left = (resized_img.width - target_size[0]) / 2
            top = resized_img.height - target_size[1]
        elif crop_mode == 'left':
            left = 0
            top = (resized_img.height - target_size[1]) / 2
        elif crop_mode == 'right':
            left = resized_img.width - target_size[0]
            top = (resized_img.height - target_size[1]) / 2

        right = left + target_size[0]
        bottom = top + target_size[1]

        # Crop the image
        cropped_img = resized_img.crop((left, top, right, bottom))

        return cropped_img

def remove_background(input_path):
    with open(input_path, 'rb') as i:
        input_image = i.read()
    output_image = remove(input_image)
    img = Image.open(io.BytesIO(output_image)).convert("RGBA")
    return img

def process_single_image(image_path, output_folder, crop_mode, remove_bg):
    filename = os.path.basename(image_path)
    try:
        if remove_bg == 'yes':
            # Remove background
            image_with_no_bg = remove_background(image_path)
            temp_image_path = os.path.join(output_folder, f"temp_{filename}")
            image_with_no_bg.save(temp_image_path, format='PNG')
        else:
            temp_image_path = image_path

        # Resize and crop the image with or without background removal
        new_image = resize_and_crop_image(temp_image_path, crop_mode=crop_mode)

        # Save the final image
        output_path = os.path.join(output_folder, f"{os.path.splitext(filename)[0]}.png")
        new_image.save(output_path, format='PNG')

        if remove_bg == 'yes':
            # Remove the temporary file
            os.remove(temp_image_path)

        return output_path

    except Exception as e:
        print(f"Error processing {filename}: {e}")
        return None

def process_images(zip_file, crop_mode='center', remove_bg='yes', progress=gr.Progress()):
    # Create a temporary directory
    input_folder = "temp_input"
    output_folder = "temp_output"
    if os.path.exists(input_folder):
        shutil.rmtree(input_folder)
    if os.path.exists(output_folder):
        shutil.rmtree(output_folder)
    os.makedirs(input_folder)
    os.makedirs(output_folder)
    
    # Extract the zip file
    with zipfile.ZipFile(zip_file, 'r') as zip_ref:
        zip_ref.extractall(input_folder)
    
    processed_images = []
    image_files = [os.path.join(input_folder, f) for f in os.listdir(input_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif'))]
    total_images = len(image_files)

    # Process images using ThreadPoolExecutor
    with ThreadPoolExecutor(max_workers=2) as executor:
        future_to_image = {executor.submit(process_single_image, image_path, output_folder, crop_mode, remove_bg): image_path for image_path in image_files}
        for idx, future in enumerate(future_to_image):
            result = future.result()
            if result:
                processed_images.append(result)
            # Update progress
            progress((idx + 1) / total_images, f"{idx + 1}/{total_images} images processed")
    
    # Create a zip file of the processed images
    output_zip_path = "processed_images.zip"
    with zipfile.ZipFile(output_zip_path, 'w') as zipf:
        for file in processed_images:
            zipf.write(file, os.path.basename(file))
    
    # Return the images and the zip file path
    return processed_images, output_zip_path

def gradio_interface(zip_file, crop_mode, remove_bg):
    progress = gr.Progress()  # Initialize progress
    return process_images(zip_file.name, crop_mode, remove_bg, progress)

# Create the Gradio interface
with gr.Blocks() as iface:
    gr.Markdown("# Image Background Removal and Resizing")
    gr.Markdown("Upload a ZIP or RAR file containing images, choose the crop mode, and get the processed images.")
    
    with gr.Row():
        zip_file = gr.File(label="Upload ZIP/RAR file of images", file_types=[".zip", ".rar"])
        crop_mode = gr.Radio(choices=["center", "top", "bottom", "left", "right"], label="Crop Mode", value="center")
        remove_bg = gr.Radio(choices=["yes", "no"], label="Remove Background", value="yes")

    gallery = gr.Gallery(label="Processed Images")
    output_zip = gr.File(label="Download Processed Images as ZIP")

    def process(zip_file, crop_mode, remove_bg):
        processed_images, zip_path = gradio_interface(zip_file, crop_mode, remove_bg)
        return processed_images, zip_path

    process_button = gr.Button("Process Images")
    process_button.click(process, inputs=[zip_file, crop_mode, remove_bg], outputs=[gallery, output_zip])

# Launch the interface
iface.launch()