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

def process_images(input_files):
    if not input_files:
        return None, "No files uploaded."

    output_folder = "processed_images"
    os.makedirs(output_folder, exist_ok=True)
    
    processed_paths = []
    
    # This logic replaces your JS loop: it renames files image1, image2...
    for i, file_path in enumerate(input_files, start=1):
        img = Image.open(file_path)
        
        # Define the new filename (e.g., image1.png)
        new_name = f"image{i}.png"
        save_path = os.path.join(output_folder, new_name)
        
        # Save the image to the new path
        img.save(save_path)
        processed_paths.append(save_path)

    # Create the ZIP file
    zip_path = "renamed_images.zip"
    with zipfile.ZipFile(zip_path, 'w') as zipf:
        for file in processed_paths:
            zipf.write(file, arcname=os.path.basename(file))
            
    return zip_path

# --- Gradio UI Design ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🖼️ Image Batch Processor")
    gr.Markdown("Upload images to rename them sequentially and download a ZIP file. [View GitHub](https://github.com)")
    
    with gr.Row():
        with gr.Column():
            # Dropzone for multiple images
            file_input = gr.File(
                label="Upload Images", 
                file_count="multiple", 
                file_types=["image"]
            )
            process_btn = gr.Button("Rewrite & Zip Files", variant="primary")
            
        with gr.Column():
            # This appears when processing is done
            file_output = gr.File(label="Download Processed ZIP")

    # Connect the logic
    process_btn.click(
        fn=process_images, 
        inputs=file_input, 
        outputs=file_output
    )

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