import os

import folder_paths
import numpy as np
import torch
from PIL import Image, ImageOps


class SaveCaptionsAndImages:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "captions": ("STRING", {"multiline": True}),
                "images": ("IMAGE",),
                "directory_prefix": (
                    "STRING",
                    {"default": "lora_dataset", "multiline": False},
                ),
            },
        }

    RETURN_TYPES = ()
    OUTPUT_NODE = True
    FUNCTION = "apply"

    def apply(self, captions, images, directory_prefix):

        # Split the captions string into a list using " | " as the delimiter
        caption_list = captions.split(" | ")
        full_output_folder = folder_paths.get_output_directory()
        # Find the next available directory number
        i = 0
        while True:
            dir_path = os.path.join(full_output_folder, f"{directory_prefix}_{i:03d}")
            if not os.path.exists(dir_path):
                break
            i += 1
        # Validate input
        if len(caption_list) != len(images):
            raise ValueError(
                "The number of captions does not match the number of images."
            )

        for batch_number, (image, caption) in enumerate(zip(images, caption_list)):
            # Generate a unique filename for each image
            filename = f"image_{batch_number:04d}"

            # Generate file paths
            image_filepath = os.path.join(dir_path, f"{filename}.png")
            caption_filepath = os.path.join(dir_path, f"{filename}.txt")

            # Ensure directory exists
            os.makedirs(dir_path, exist_ok=True)

            # Save the image
            i = 255.0 * image.cpu().numpy()
            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
            img.save(image_filepath)

            # Write caption to file
            with open(caption_filepath, "w", encoding="utf-8") as caption_file:
                caption_file.write(caption)

            print(f"Image saved to: {image_filepath}")
            print(f"Caption saved to: {caption_filepath}")

        return {}


class BizyAirLoadImagesFromFolder:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "folder": ("STRING", {"default": ""}),
                "width": ("INT", {"default": 1024, "min": 64, "step": 1}),
                "height": ("INT", {"default": 1024, "min": 64, "step": 1}),
                "keep_aspect_ratio": (
                    [
                        "crop",
                        "pad",
                        "stretch",
                    ],
                ),
            },
            "optional": {
                "image_load_cap": ("INT", {"default": 100, "min": 0, "step": 1}),
                "start_index": ("INT", {"default": 0, "min": 0, "step": 1}),
                "include_subfolders": ("BOOLEAN", {"default": False}),
            },
        }

    RETURN_TYPES = (
        "IMAGE",
        "MASK",
        "INT",
        "STRING",
    )
    RETURN_NAMES = (
        "image",
        "mask",
        "count",
        "image_path",
    )
    FUNCTION = "load_images"
    CATEGORY = "☁️BizyAir/marking"
    DESCRIPTION = """Loads images from a folder into a batch, images are resized and loaded into a batch."""

    def load_images(
        self,
        folder,
        width,
        height,
        image_load_cap,
        start_index,
        keep_aspect_ratio,
        include_subfolders=False,
    ):
        if not os.path.isdir(folder):
            raise FileNotFoundError(f"Folder '{folder} cannot be found.'")

        valid_extensions = [".jpg", ".jpeg", ".png", ".webp"]
        image_paths = []
        if include_subfolders:
            for root, _, files in os.walk(folder):
                for file in files:
                    if any(file.lower().endswith(ext) for ext in valid_extensions):
                        image_paths.append(os.path.join(root, file))
        else:
            for file in os.listdir(folder):
                if any(file.lower().endswith(ext) for ext in valid_extensions):
                    image_paths.append(os.path.join(folder, file))

        dir_files = sorted(image_paths)

        if len(dir_files) == 0:
            raise FileNotFoundError(f"No files in directory '{folder}'.")

        # start at start_index
        dir_files = dir_files[start_index:]

        images = []
        masks = []
        image_path_list = []

        limit_images = False
        if image_load_cap > 0:
            limit_images = True
        image_count = 0

        for image_path in dir_files:
            if os.path.isdir(image_path):
                continue
            if limit_images and image_count >= image_load_cap:
                break
            i = Image.open(image_path)
            i = ImageOps.exif_transpose(i)

            # Resize image to maximum dimensions
            if i.size != (width, height):
                i = self.resize_with_aspect_ratio(i, width, height, keep_aspect_ratio)

            image = i.convert("RGB")
            image = np.array(image).astype(np.float32) / 255.0
            image = torch.from_numpy(image)[None,]

            if "A" in i.getbands():
                mask = np.array(i.getchannel("A")).astype(np.float32) / 255.0
                mask = 1.0 - torch.from_numpy(mask)
                if mask.shape != (height, width):
                    mask = torch.nn.functional.interpolate(
                        mask.unsqueeze(0).unsqueeze(0),
                        size=(height, width),
                        mode="bilinear",
                        align_corners=False,
                    ).squeeze()
            else:
                mask = torch.zeros((height, width), dtype=torch.float32, device="cpu")

            images.append(image)
            masks.append(mask)
            image_path_list.append(image_path)
            image_count += 1

        if len(images) == 1:
            return (images[0], masks[0], 1, image_path_list)

        elif len(images) > 1:
            image1 = images[0]
            mask1 = masks[0].unsqueeze(0)

            for image2 in images[1:]:
                image1 = torch.cat((image1, image2), dim=0)

            for mask2 in masks[1:]:
                mask1 = torch.cat((mask1, mask2.unsqueeze(0)), dim=0)

            return (image1, mask1, len(images), image_path_list)

    def resize_with_aspect_ratio(self, img, width, height, mode):
        if mode == "stretch":
            return img.resize((width, height), Image.Resampling.LANCZOS)

        img_width, img_height = img.size
        aspect_ratio = img_width / img_height
        target_ratio = width / height

        if mode == "crop":
            # Calculate dimensions for center crop
            if aspect_ratio > target_ratio:
                # Image is wider - crop width
                new_width = int(height * aspect_ratio)
                img = img.resize((new_width, height), Image.Resampling.LANCZOS)
                left = (new_width - width) // 2
                return img.crop((left, 0, left + width, height))
            else:
                # Image is taller - crop height
                new_height = int(width / aspect_ratio)
                img = img.resize((width, new_height), Image.Resampling.LANCZOS)
                top = (new_height - height) // 2
                return img.crop((0, top, width, top + height))

        elif mode == "pad":
            pad_color = self.get_edge_color(img)
            # Calculate dimensions for padding
            if aspect_ratio > target_ratio:
                # Image is wider - pad height
                new_height = int(width / aspect_ratio)
                img = img.resize((width, new_height), Image.Resampling.LANCZOS)
                padding = (height - new_height) // 2
                padded = Image.new("RGBA", (width, height), pad_color)
                padded.paste(img, (0, padding))
                return padded
            else:
                # Image is taller - pad width
                new_width = int(height * aspect_ratio)
                img = img.resize((new_width, height), Image.Resampling.LANCZOS)
                padding = (width - new_width) // 2
                padded = Image.new("RGBA", (width, height), pad_color)
                padded.paste(img, (padding, 0))
                return padded

    def get_edge_color(self, img):
        from PIL import ImageStat

        """Sample edges and return dominant color"""
        width, height = img.size
        img = img.convert("RGBA")

        # Create 1-pixel high/wide images from edges
        top = img.crop((0, 0, width, 1))
        bottom = img.crop((0, height - 1, width, height))
        left = img.crop((0, 0, 1, height))
        right = img.crop((width - 1, 0, width, height))

        # Combine edges into single image
        edges = Image.new("RGBA", (width * 2 + height * 2, 1))
        edges.paste(top, (0, 0))
        edges.paste(bottom, (width, 0))
        edges.paste(left.resize((height, 1)), (width * 2, 0))
        edges.paste(right.resize((height, 1)), (width * 2 + height, 0))

        # Get median color
        stat = ImageStat.Stat(edges)
        median = tuple(map(int, stat.median))
        return median


NODE_CLASS_MAPPINGS = {
    "BizyAirLoadImagesFromFolder": BizyAirLoadImagesFromFolder,
    "SaveCaptionsAndImages": SaveCaptionsAndImages,
}
NODE_DISPLAY_NAME_MAPPINGS = {
    "BizyAirLoadImagesFromFolder": "☁️BizyAir LoadImagesFromFolder",
    "SaveCaptionsAndImages": "☁️BizyAir Save Captions And Images",
}