from processing import process_images
import gradio as gr
from PIL import Image
import os
import requests
from io import BytesIO




def gradio_interface(input_files, bg_method, watermark, canvas_size, output_format, bg_choice, custom_color, num_workers):
    progress = gr.Progress()
    watermark_path = watermark.name if watermark else None
    
    # Check input_files, is it single image, list image, or zip/rar
    if isinstance(input_files, str) and input_files.lower().endswith(('.zip', '.rar')):
            return process_images(input_files, bg_method, watermark_path, canvas_size, output_format, bg_choice, custom_color, num_workers, progress)
    elif isinstance(input_files, list):
        return process_images(input_files, bg_method, watermark_path, canvas_size, output_format, bg_choice, custom_color, num_workers, progress)
    else:
        return process_images(input_files.name, bg_method, watermark_path, canvas_size, 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_compare(evt: gr.SelectData):
    print(f"Selected value: {evt.value}")  # Debug print
    
    try:
        if isinstance(evt.value, list) and len(evt.value) == 2:
            _, combined_path = evt.value
        elif isinstance(evt.value, str):
            combined_path = evt.value
        else:
            raise ValueError(f"Unexpected input format: {evt.value}")
        
        output_path, input_path = combined_path.split('|')
        
        # Handle URLs for deployed version
        if output_path.startswith('http'):
            output_img = Image.open(BytesIO(requests.get(output_path).content))
        else:
            output_path = os.path.join("processed_images", os.path.basename(output_path))
            output_img = Image.open(output_path)
        
        if input_path.startswith('http'):
            input_img = Image.open(BytesIO(requests.get(input_path).content))
        else:
            input_path = os.path.join("temp_input", os.path.basename(input_path))
            input_img = Image.open(input_path)
        
        # Calculate the aspect ratios
        original_ratio = f"{input_img.width}x{input_img.height}"
        processed_ratio = f"{output_img.width}x{output_img.height}"
        
        print(f"Successfully processed. Input: {input_path}, Output: {output_path}")
        return gr.update(value=input_img), gr.update(value=output_img), gr.update(value=original_ratio), gr.update(value=processed_ratio)
    except Exception as e:
        print(f"Error in update_compare: {e}")
        return gr.update(value=None), gr.update(value=None), gr.update(value=None), gr.update(value=None)


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

with gr.Blocks(theme="NoCrypt/miku@1.2.2") 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.")

    with gr.Row():
        input_files = gr.File(label="Upload Image or ZIP/RAR file", file_types=[".zip", ".rar", "image"], interactive=True)
        watermark = gr.File(label="Upload Watermark Image (Optional)", file_types=[".png"])

    with gr.Row():
        canvas_size = gr.Radio(choices=["Rox", "Columbia", "Zalora"], label="Canvas Size", value="Rox")
        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", "none"], label="Background Removal Method", value="bria")
        bg_choice = gr.Radio(choices=["transparent", "white", "custom", "blur"], label="Background Choice", value="transparent")
        custom_color = gr.ColorPicker(label="Custom Background Color", value="#ffffff", visible=False)

    process_button = gr.Button("Process Images")

    with gr.Row():
        gallery_processed = gr.Gallery(label="Processed Images", show_label=True, elem_id="gallery")
    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)")

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

    
    process_button.click(process, inputs=[input_files, bg_method, watermark, canvas_size, output_format, bg_choice, custom_color, num_workers], outputs=[gallery_processed, output_zip, processing_time])
    gallery_processed.select(update_compare, outputs=[image_original, image_processed, original_ratio, processed_ratio])

if __name__ == "__main__":
    iface.launch(share=True)