import gradio as gr
import numpy as np
from PIL import Image
import tempfile
import shutil
import os
from zipfile import ZipFile
import io

def sepia(input_img, num_copies):
    sepia_filter = np.array([
        [0.393, 0.769, 0.189], 
        [0.349, 0.686, 0.168], 
        [0.272, 0.534, 0.131]
    ])
    input_img = np.array(input_img) / 255.0
    sepia_imgs = []
    for _ in range(num_copies):
        sepia_img = np.dot(input_img[..., :3], sepia_filter.T)
        sepia_img = np.clip(sepia_img, 0, 1) * 255
        sepia_imgs.append(Image.fromarray(sepia_img.astype(np.uint8)))
    return sepia_imgs


def zip_sepia_images(sepia_imgs):
    # Initialize a BytesIO object to hold the ZIP file in memory
    zip_bytes = io.BytesIO()
    with ZipFile(zip_bytes, 'w') as zipf:
        for i, img in enumerate(sepia_imgs):
            # Convert the PIL Image to a bytes object
            img_byte_arr = io.BytesIO()
            img.save(img_byte_arr, format='PNG')
            img_byte_arr.seek(0)
            # Write the image bytes to the ZIP file
            zipf.writestr(f"sepia_image_{i}.png", img_byte_arr.getvalue())
    # IMPORTANT: Seek back to the start of the BytesIO object
    zip_bytes.seek(0)
    
    # Return the ZIP file as a tuple: (filename, bytes)
    # Note: Ensure the filename has the correct file extension (.zip)
    return "sepia_images.zip", zip_bytes.getvalue()


with gr.Blocks() as demo:
    with gr.Row():
        input_img = gr.Image()
        num_copies = gr.Number(label="Number of Copies", value=1)
    with gr.Row():
        gallery = gr.Gallery(label="Sepia Images")
        download_btn = gr.File(label="Download ZIP")

def update_output(input_img, num_copies):
    sepia_imgs = sepia(input_img, num_copies)
    zip_name, zip_bytes = zip_sepia_images(sepia_imgs)
    return sepia_imgs, (zip_name, zip_bytes)  # This tuple matches Gradio's expectation

    generate_btn = gr.Button("Generate Sepia Images")
    generate_btn.click(fn=update_output, 
                       inputs=[input_img, num_copies], 
                       outputs=[gallery, download_btn])

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