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 title(self):
        return "Image Size Randomize"
      
    def ui(self, is_txt2img):
        with gr.Tabs():
            # Tab for direct text input
            with gr.TabItem("Image Sizes"):
                gr.HTML("<p>Enter image sizes in the format 'width,height', separated by new lines.</p>")
                with gr.Row():
                    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  # Adjust line count for better usability
                    )
                with gr.Row():
                    random_swap_text = gr.Checkbox(
                        label='Randomly swap width and height',
                        elem_id=self.elem_id("random_swap_text")
                    )

        return [text_sizes, random_swap_text]

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

        # Read sizes from the textbox
        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 in textbox: {text_sizes}")

        if sizes:
            width, height = random.choice(sizes)
            if random_swap_text:
                if random.choice([True, False]):
                    width, height = height, width
            p.width, p.height = width, height
        else:
            print(f"No valid sizes found. Please check the textbox input.")

        proc = process_images(p)
        return proc