import random
import tempfile
import zipfile
import os
import uuid
import PIL.Image
import PIL.ImageOps
import gradio as gr

from settings import MAX_SEED


def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)  # noqa: S311
    return seed


def wrap_process_with_zip(process_fn, has_color_correction=True):
    def wrapped(*args):
        try:
            if has_color_correction:
                # The last argument of args is the color_correction value
                color_corr_val = args[-1]
                # The remaining arguments are passed to the original process_fn
                actual_args = args[:-1]
            else:
                color_corr_val = None
                actual_args = args

            res_images = process_fn(*actual_args)
        except Exception as e:
            print(f"Error in process: {e}")
            raise e

        if not res_images:
            return [], gr.update(visible=False)

        processed_res = []
        for img in res_images:
            if color_corr_val == "white bg, black lines":
                # Invert the image (black bg -> white bg, white lines -> black lines)
                if isinstance(img, PIL.Image.Image):
                    img_rgb = img.convert("RGB")
                    inverted_img = PIL.ImageOps.invert(img_rgb)
                    processed_res.append(inverted_img)
                elif isinstance(img, str) and os.path.exists(img):
                    try:
                        img_pil = PIL.Image.open(img).convert("RGB")
                        inverted_img = PIL.ImageOps.invert(img_pil)
                        processed_res.append(inverted_img)
                    except Exception:
                        processed_res.append(img)
                else:
                    # Try to convert numpy array or other format
                    try:
                        img_pil = PIL.Image.fromarray(img).convert("RGB")
                        inverted_img = PIL.ImageOps.invert(img_pil)
                        processed_res.append(inverted_img)
                    except Exception:
                        processed_res.append(img)
            else:
                processed_res.append(img)

        res_images = processed_res

        # Now create the ZIP file from res_images
        temp_dir = tempfile.gettempdir()
        zip_filename = os.path.join(temp_dir, f"controlnet_results_{uuid.uuid4().hex[:8]}.zip")

        try:
            with zipfile.ZipFile(zip_filename, 'w') as zipf:
                for idx, img in enumerate(res_images):
                    if isinstance(img, PIL.Image.Image):
                        temp_img_path = os.path.join(temp_dir, f"temp_{uuid.uuid4().hex[:8]}.png")
                        img.save(temp_img_path, "PNG")
                        zipf.write(temp_img_path, f"result_{idx}.png")
                        try:
                            os.remove(temp_img_path)
                        except Exception:
                            pass
                    elif isinstance(img, str) and os.path.exists(img):
                        zipf.write(img, os.path.basename(img))
                    else:
                        try:
                            temp_img_path = os.path.join(temp_dir, f"temp_{uuid.uuid4().hex[:8]}.png")
                            PIL.Image.fromarray(img).save(temp_img_path, "PNG")
                            zipf.write(temp_img_path, f"result_{idx}.png")
                            try:
                                os.remove(temp_img_path)
                            except Exception:
                                pass
                        except Exception:
                            pass
            return res_images, gr.update(value=zip_filename, visible=True)
        except Exception as zip_err:
            print(f"Error creating zip file: {zip_err}")
            return res_images, gr.update(visible=False)

    return wrapped