import random
import modules.scripts as scripts
import gradio as gr
import os

from modules import images
from modules.processing import process_images, Processed
from modules.shared import opts, cmd_opts, state

class RandomImageSizeScript(scripts.Script):
    def __init__(self):
        super().__init__()
        self.index = 0  # Round-Robin mode index
        self.swap_round = False  # Controls width-height swap per round

    def title(self):
        return "Image Size Randomize"

    def ui(self, is_txt2img):
        with gr.Accordion("Image Size Randomize Settings", open=True):
            with gr.Group():
                gr.Markdown("### Size")
                gr.HTML("<p>Enter image sizes in the format 'width,height', separated by new lines.</p>")
                text_sizes = gr.Textbox(
                    label='Image Sizes',
                    placeholder='Enter sizes, e.g.:\n1024,512\n1024,768\n1024,1024',
                    elem_id=self.elem_id("text_sizes"),
                    lines=10
                )

            with gr.Group():
                gr.Markdown("### Mode")
                mode = gr.Radio(
                    choices=["Random", "Round-Robin"],
                    value="Random",
                    label="Selection Mode",
                    elem_id=self.elem_id("mode"),
                    info="Choose whether to select sizes randomly or in a round-robin fashion."
                )

                random_swap_text = gr.Checkbox(
                    label='Randomly swap width and height',
                    elem_id=self.elem_id("random_swap_text"),
                    visible=True
                )
                
                sequential_swap = gr.Checkbox(
                    label='Swap width and height every round',
                    elem_id=self.elem_id("sequential_swap"),
                    visible=False
                )

                def update_visibility(selected_mode):
                    if selected_mode == "Random":
                        return {
                            random_swap_text: gr.update(visible=True),
                            sequential_swap: gr.update(visible=False)
                        }
                    else:
                        return {
                            random_swap_text: gr.update(visible=False),
                            sequential_swap: gr.update(visible=True)
                        }

                mode.change(
                    fn=update_visibility,
                    inputs=mode,
                    outputs=[random_swap_text, sequential_swap]
                )

        return [text_sizes, random_swap_text, mode, sequential_swap]

    def run(self, p, text_sizes, random_swap_text, mode, sequential_swap):
        sizes = []

        # Parse and validate input sizes
        if text_sizes:
            try:
                sizes = [tuple(map(int, line.strip().split(','))) for line in text_sizes.strip().split('\n') if line.strip()]
            except ValueError:
                print(f"Invalid size format detected. Ensure all sizes are in the format 'width,height'.")
        
        if sizes:
            try:
                if mode == "Random":
                    width, height = random.choice(sizes)
                    if random_swap_text and random.choice([True, False]):
                        width, height = height, width
                else:  # Round-Robin mode
                    width, height = sizes[self.index]

                    # Handle dynamic size changes and index resets
                    if self.index >= len(sizes):
                        self.index = 0  # Reset index if out of bounds
                        width, height = sizes[self.index]

                    if self.index == 0 and sequential_swap:
                        self.swap_round = not self.swap_round

                    if sequential_swap and self.swap_round:
                        width, height = height, width

                    # Update index for the next round
                    self.index = (self.index + 1) % len(sizes)

                # Set dimensions for image processing
                p.width, p.height = width, height

            except IndexError as e:
                print(f"Error: {e}. The size list may have been modified unexpectedly.")
                self.index = 0  # Reset index and recover from error

        else:
            print("No valid sizes found. Please check the input format.")

        proc = process_images(p)
        return proc