import os.path
import modules.scripts as scripts
import gradio as gr
from modules import sd_samplers, shared
from modules.processing import Processed, process_images, StableDiffusionProcessing, create_infotext
import modules.images as images
from modules.shared import opts, cmd_opts, state
from PIL import Image
import os

class Script(scripts.Script):
    def title(self):
        return "GIF creator from Image Slice (multiple)"

    def show(self, is_txt2img):
        return True

    def ui(self, is_txt2img):
        num_cuts = gr.inputs.Slider(minimum=2, maximum=8, default=2, step=1, label="Number of cuts/slices")
        gif_duration = gr.inputs.Slider(minimum=5, maximum=1000, default=150, step=5, label="GIF duration (ms)")
        ping_pong_checkbox = gr.inputs.Checkbox(label="Ping-pong animation", default=True)
        return [num_cuts, gif_duration, ping_pong_checkbox]

    def cut_image(self, image, num_cuts):
        width, height = image.size
        cut_width = width // num_cuts
        cut_height = height // num_cuts
        parts = []
        for i in range(num_cuts):
            for j in range(num_cuts):
                left = j * cut_width
                upper = i * cut_height
                right = left + cut_width
                lower = upper + cut_height
                parts.append(image.crop((left, upper, right, lower)))
        return parts

    def run(self, p, num_cuts, gif_duration, ping_pong_animation):
        proc = process_images(p)
        gens = proc.images

        print("Start processing "+str(len(gens))+" gif(s)")
        for image_index in range(len(gens)):
            save_image_tuple = images.save_image(gens[image_index], p.outpath_samples, "", image_index, "-Original", opts.samples_format)
            save_image = save_image_tuple[0]
            image = Image.open(save_image) # Cut the image into equal parts and create a GIF with optional ping-pong animation
            cut_parts = self.cut_image(image, num_cuts)
            gif_filename = os.path.splitext(save_image)[0] + "_gif.gif"
            if ping_pong_animation:
                cut_parts = cut_parts + cut_parts[-2:0:-1]  # Create a ping-pong sequence
            cut_parts[0].save(gif_filename, format='GIF', append_images=cut_parts[1:], save_all=True, duration=gif_duration, loop=0)
            print(f"GIF created: {gif_filename}")
            os.remove(save_image) # Delete the original saved image

        return Processed(p, gens, 0, "")