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

def colors_within_tolerance(color1, color2, tolerance):
    return all(abs(c1 - c2) <= tolerance for c1, c2 in zip(color1, color2))

def check_border_colors(image_path, tolerance):
    image = Image.open(image_path)
    pixels = image.load()
    
    width, height = image.size

    left_border_color = pixels[0, 0]
    right_border_color = pixels[width - 1, 0]

    for y in range(height):
        if not colors_within_tolerance(pixels[0, y], left_border_color, tolerance):
            return False
        if not colors_within_tolerance(pixels[width - 1, y], right_border_color, tolerance):
            return False

    return True

def resize_and_crop_image(image_path, target_size=(1080, 1080), crop_mode='center'):
    print(f"Resizing and cropping image: {image_path}")
    with Image.open(image_path) as img:
        width, height = img.size
        print(f"Original image size: {width}x{height}")

        scaling_factor = max(target_size[0] / width, target_size[1] / height)

        new_size = (int(width * scaling_factor), int(height * scaling_factor))
        resized_img = img.resize(new_size, Image.LANCZOS)
        print(f"Resized image size: {new_size}")

        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]

        cropped_img = resized_img.crop((left, top, right, bottom))
        print(f"Cropped image size: {cropped_img.size}")

        return cropped_img

def remove_background_rembg(input_path):
    print(f"Removing background using rembg for image: {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 remove_background_bria(input_path):
    print(f"Removing background using bria for image: {input_path}")
    pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True, device=0)
    pillow_image = pipe(input_path)
    return pillow_image

def process_single_image(image_path, output_folder, crop_mode, bg_method, output_format, bg_choice, custom_color, watermark_path=None):
    filename = os.path.basename(image_path)
    try:
        print(f"Processing image: {filename}")
        if bg_method == 'rembg':
            image_with_no_bg = remove_background_rembg(image_path)
        elif bg_method == 'bria':
            image_with_no_bg = remove_background_bria(image_path)
            
        temp_image_path = os.path.join(output_folder, f"temp_{filename}")
        image_with_no_bg.save(temp_image_path, format='PNG')

        if check_border_colors(temp_image_path, tolerance=50):
            print(f"Border colors are the same for image: {filename}")
            if bg_choice == 'transparent':
                new_image = Image.new("RGBA", (1080, 1080), (255, 255, 255, 0))
            else:
                new_image = Image.new("RGBA", (1080, 1080), custom_color)

            width, height = image_with_no_bg.size
            scaling_factor = min(1080 / width, 1080 / height)
            new_size = (int(width * scaling_factor), int(height * scaling_factor))
            resized_img = image_with_no_bg.resize(new_size, Image.LANCZOS)
            print(f"Resized image size: {new_size}")
            new_image.paste(resized_img, ((1080 - resized_img.width) // 2, (1080 - resized_img.height) // 2))
        else:
            print(f"Border colors are different for image: {filename}")
            new_image = resize_and_crop_image(temp_image_path, crop_mode=crop_mode)

        if bg_choice == 'white':
            new_image = new_image.convert("RGBA")
            white_bg = Image.new("RGBA", new_image.size, "WHITE")
            new_image = Image.alpha_composite(white_bg, new_image)
        elif bg_choice == 'custom':
            new_image = new_image.convert("RGBA")
            custom_bg = Image.new("RGBA", new_image.size, custom_color)
            new_image = Image.alpha_composite(custom_bg, new_image)

        images_paths = []

        output_ext = 'jpg' if output_format == 'JPG' else 'png'
        output_path_without_watermark = os.path.join(output_folder, f"without_watermark_{os.path.splitext(filename)[0]}.{output_ext}")
        if output_format == 'JPG':
            new_image.convert('RGB').save(output_path_without_watermark, format='JPEG')
        else:
            new_image.save(output_path_without_watermark, format='PNG')
        images_paths.append((output_path_without_watermark, image_path))

        if watermark_path:
            watermark = Image.open(watermark_path).convert("RGBA")
            new_image_with_watermark = new_image.copy()
            new_image_with_watermark.paste(watermark, (0, 0), watermark)
            output_path_with_watermark = os.path.join(output_folder, f"with_watermark_{os.path.splitext(filename)[0]}.{output_ext}")
            if output_format == 'JPG':
                new_image_with_watermark.convert('RGB').save(output_path_with_watermark, format='JPEG')
            else:
                new_image_with_watermark.save(output_path_with_watermark, format='PNG')
            images_paths.append((output_path_with_watermark, image_path))

        os.remove(temp_image_path)

        print(f"Processed image paths: {images_paths}")
        return images_paths

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

def process_images(input_files, crop_mode='center', bg_method='rembg', watermark_path=None, output_format='PNG', bg_choice='transparent', custom_color="#ffffff", num_workers=4, progress=gr.Progress()):
    start_time = time.time()
    
    output_folder = "temp_output"
    if os.path.exists(output_folder):
        shutil.rmtree(output_folder)
    os.makedirs(output_folder)
    
    processed_images = []
    original_images = []
    
    if isinstance(input_files, str) and input_files.lower().endswith(('.zip', '.rar')):
        input_folder = "temp_input"
        if os.path.exists(input_folder):
            shutil.rmtree(input_folder)
        os.makedirs(input_folder)
        
        try:
            with zipfile.ZipFile(input_files, 'r') as zip_ref:
                zip_ref.extractall(input_folder)
        except zipfile.BadZipFile as e:
            print(f"Error extracting zip file: {e}")
            return [], None, 0
        
        image_files = [os.path.join(input_folder, f) for f in os.listdir(input_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif', '.webp'))]
    else:
        image_files = [f.name for f in input_files]
    
    total_images = len(image_files)
    print(f"Total images to process: {total_images}")

    avg_processing_time = 0
    with ThreadPoolExecutor(max_workers=num_workers) as executor:
        future_to_image = {executor.submit(process_single_image, image_path, output_folder, crop_mode, bg_method, output_format, bg_choice, custom_color, watermark_path): image_path for image_path in image_files}
        for idx, future in enumerate(future_to_image):
            try:
                start_time_image = time.time()
                result = future.result()
                end_time_image = time.time()
                image_processing_time = end_time_image - start_time_image
                
                # Update average processing time
                avg_processing_time = (avg_processing_time * idx + image_processing_time) / (idx + 1)
                
                if result:
                    processed_images.extend(result)
                    original_images.append(future_to_image[future])
                
                # Estimate remaining time
                remaining_images = total_images - (idx + 1)
                estimated_remaining_time = remaining_images * avg_processing_time
                
                progress((idx + 1) / total_images, f"{idx + 1}/{total_images} images processed. Estimated time remaining: {estimated_remaining_time:.2f} seconds")
            except Exception as e:
                print(f"Error processing image {future_to_image[future]}: {e}")

    output_zip_path = "processed_images.zip"
    with zipfile.ZipFile(output_zip_path, 'w') as zipf:
        for file, _ in processed_images:
            if "with_watermark" in file:
                zipf.write(file, os.path.join("with_watermark", os.path.basename(file)))
            else:
                zipf.write(file, os.path.join("without_watermark", os.path.basename(file)))

    end_time = time.time()
    processing_time = end_time - start_time
    print(f"Processing time: {processing_time} seconds")
    
    return original_images, processed_images, output_zip_path, processing_time

def gradio_interface(input_type, input_files, crop_mode, bg_method, watermark, output_format, bg_choice, custom_color, num_workers):
    progress = gr.Progress()
    watermark_path = watermark.name if watermark else None
    
    if input_type == "zip_rar":
        return process_images(input_files.name, crop_mode, bg_method, watermark_path, output_format, bg_choice, custom_color, num_workers, progress)
    else:
        return process_images(input_files, crop_mode, bg_method, watermark_path, output_format, bg_choice, custom_color, num_workers, progress)

def show_color_picker(bg_choice):
    if bg_choice == 'custom':
        return gr.update(visible=True)
    return gr.update(visible=False)

def update_slider(evt: gr.SelectData):
    if isinstance(evt.value, dict) and 'caption' in evt.value:
        input_path = evt.value['caption']
        output_path = evt.value['image']['path']
        input_path = input_path.split("Input: ")[-1]
        
        # Open the original and processed images
        original_img = Image.open(input_path)
        processed_img = Image.open(output_path)
        
        # Calculate the aspect ratios
        original_ratio = f"{original_img.width}x{original_img.height}"
        processed_ratio = f"{processed_img.width}x{processed_img.height}"
        
        return gr.update(value=input_path), gr.update(value=output_path), gr.update(value=original_ratio), gr.update(value=processed_ratio)
    else:
        print("No caption found in selection")
        return gr.update(value=None), gr.update(value=None), gr.update(value=None), gr.update(value=None)

def process(input_type, input_files, crop_mode, bg_method, watermark, output_format, bg_choice, custom_color, num_workers):
    _, processed_images, zip_path, time_taken = gradio_interface(input_type, input_files, crop_mode, bg_method, watermark, output_format, bg_choice, custom_color, num_workers)
    processed_images_with_captions = [(img, f"Input: {caption}") for img, caption in processed_images]
    return processed_images_with_captions, zip_path, f"{time_taken:.2f} seconds"

def update_input_type(choice):
    if choice == "zip_rar":
        return gr.update(visible=True), gr.update(visible=False)
    else:
        return gr.update(visible=False), gr.update(visible=True)

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

    input_type = gr.Radio(choices=["multiple_images", "zip_rar"], label="Input Type", value="multiple_images")
    
    with gr.Row():
        zip_file = gr.File(label="Upload ZIP/RAR file of images", file_types=[".zip", ".rar"], visible=False)
        multiple_images = gr.File(label="Upload multiple images", file_types=["image"], file_count="multiple")
    
    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")
        output_format = gr.Radio(choices=["PNG", "JPG"], label="Output Format", value="JPG")
        num_workers = gr.Slider(minimum=1, maximum=16, step=1, label="Number of Workers", value=5)

    with gr.Row():
        bg_method = gr.Radio(choices=["bria", "rembg"], label="Background Removal Method", value="bria")
        bg_choice = gr.Radio(choices=["transparent", "white", "custom"], label="Background Choice", value="white")
        custom_color = gr.ColorPicker(label="Custom Background Color", value="#ffffff", visible=False)

    with gr.Row():
        gallery_processed = gr.Gallery(label="Processed Images")
    with gr.Row():
        image_original = gr.Image(label="Original Images", interactive=False)
        image_processed = gr.Image(label="Processed Images", interactive=False)
    with gr.Row():
        original_ratio = gr.Textbox(label="Original Ratio")
        processed_ratio = gr.Textbox(label="Processed Ratio")
    with gr.Row():
        output_zip = gr.File(label="Download Processed Images as ZIP")
        processing_time = gr.Textbox(label="Processing Time (seconds)")

    input_type.change(update_input_type, inputs=input_type, outputs=[zip_file, multiple_images])

    bg_choice.change(show_color_picker, inputs=bg_choice, outputs=custom_color)

    process_button = gr.Button("Process Images")
    process_button.click(process, inputs=[input_type, multiple_images, crop_mode, bg_method, watermark, output_format, bg_choice, custom_color, num_workers], outputs=[gallery_processed, output_zip, processing_time])

    gallery_processed.select(update_slider, outputs=[image_original, image_processed, original_ratio, processed_ratio])

iface.launch()