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, output_format, bg_choice, watermark_path=None):
    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)

        # Apply watermark if provided
        if watermark_path:
            watermark = Image.open(watermark_path).convert("RGBA")
            new_image.paste(watermark, (0, 0), watermark)

        # Save the final image
        if remove_bg == 'yes' and output_format == 'PNG' and bg_choice == 'transparent':
            output_path = os.path.join(output_folder, f"{os.path.splitext(filename)[0]}.png")
            new_image.save(output_path, format='PNG')
        elif remove_bg == 'yes' and output_format == 'PNG' and bg_choice == 'white':
            # Create a white background
            background = Image.new("RGBA", new_image.size, (255, 255, 255, 255))
            # Paste the image with no background onto the white background
            background.paste(new_image, mask=new_image.split()[3])  # 3 is the alpha channel
            output_path = os.path.join(output_folder, f"{os.path.splitext(filename)[0]}.png")
            background.convert('RGB').save(output_path, format='PNG')
        elif remove_bg == 'yes' and output_format == 'JPG':
            # Create a white background
            background = Image.new("RGB", new_image.size, (255, 255, 255))
            # Paste the image with no background onto the white background
            background.paste(new_image, mask=new_image.split()[3])  # 3 is the alpha channel
            output_path = os.path.join(output_folder, f"{os.path.splitext(filename)[0]}.jpg")
            background.save(output_path, format='JPEG')
        else:
            output_ext = 'jpg' if output_format == 'JPG' else 'png'
            output_path = os.path.join(output_folder, f"{os.path.splitext(filename)[0]}.{output_ext}")
            if output_format == 'JPG':
                new_image.convert('RGB').save(output_path, format='JPEG')
            else:
                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', watermark_path=None, output_format='PNG', bg_choice='transparent', 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, output_format, bg_choice, watermark_path): 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, watermark, output_format, bg_choice):
    progress = gr.Progress()  # Initialize progress
    watermark_path = watermark.name if watermark else None
    return process_images(zip_file.name, crop_mode, remove_bg, watermark_path, output_format, bg_choice, progress)

def show_bg_choice(remove_bg, output_format):
    if remove_bg == 'yes' and output_format == 'PNG':
        return gr.update(visible=True)
    return gr.update(visible=False)

# Create the Gradio interface
with gr.Blocks() as iface:
    gr.Markdown("# Image Background Removal and Resizing with Optional Watermark")
    gr.Markdown("Upload a ZIP or RAR file containing images, choose the crop mode, optionally upload a watermark image, and select the output format.")

    with gr.Row():
        zip_file = gr.File(label="Upload ZIP/RAR file of images", file_types=[".zip", ".rar"])
        watermark = gr.File(label="Upload Watermark Image (Optional)", file_types=[".png"])
    
    with gr.Row():
        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")
        output_format = gr.Radio(choices=["PNG", "JPG"], label="Output Format", value="PNG")
        bg_choice = gr.Radio(choices=["transparent", "white"], label="Background Choice", value="transparent", visible=True)

    remove_bg.change(show_bg_choice, inputs=[remove_bg, output_format], outputs=bg_choice)
    output_format.change(show_bg_choice, inputs=[remove_bg, output_format], outputs=bg_choice)

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

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

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

# Launch the interface
iface.launch()