import gradio as gr
import os
import zipfile
from PIL import Image

def process_and_preview(input_files):
    if not input_files:
        return None, None
    
    output_folder = "processed_images"
    os.makedirs(output_folder, exist_ok=True)
    
    processed_paths = []
    
    # Process images
    for i, file_path in enumerate(input_files, start=1):
        img = Image.open(file_path)
        new_name = f"image{i}.png"
        save_path = os.path.join(output_folder, new_name)
        img.save(save_path)
        processed_paths.append(save_path)

    # Create ZIP
    zip_path = "processed_batch.zip"
    with zipfile.ZipFile(zip_path, 'w') as zipf:
        for file in processed_paths:
            zipf.write(file, arcname=os.path.basename(file))
            
    # Return both the zip file and the list of paths for the Gallery preview
    return zip_path, processed_paths

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🖼️ Image Batch Processor (Gradio 6.9)")
    gr.Markdown("Upload images, rename them, and download the ZIP. [GitHub](https://github.com)")
    
    with gr.Row():
        with gr.Column(scale=1):
            file_input = gr.File(
                label="Step 1: Upload Images", 
                file_count="multiple", 
                file_types=["image"]
            )
            process_btn = gr.Button("Step 2: Process & Zip", variant="primary")
            
        with gr.Column(scale=2):
            # The preview gallery
            preview_gallery = gr.Gallery(
                label="Image Preview", 
                columns=4, 
                height="auto",
                object_fit="contain"
            )
            # The final download button/file
            zip_output = gr.File(label="Step 3: Download ZIP")

    # Connect the button to the function
    process_btn.click(
        fn=process_and_preview, 
        inputs=file_input, 
        outputs=[zip_output, preview_gallery]
    )

if __name__ == "__main__":
    demo.launch()