"""
Example ComfyUI node for pixel sorting.

This is a template showing how to integrate pixelsort into ComfyUI.
Place this file in your ComfyUI custom_nodes directory.
"""

import torch
import numpy as np
from PIL import Image

from pixelsort import pixelsort, analyze, INTERVAL_FUNCTIONS, SORTING_FUNCTIONS


class PixelSortNode:
    """ComfyUI node for pixel sorting effects."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "interval": (list(INTERVAL_FUNCTIONS.keys()), {"default": "bright"}),
                "sorting": (list(SORTING_FUNCTIONS.keys()), {"default": "value"}),
            },
            "optional": {
                "angle": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 1.0}),
                "randomness": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100.0, "step": 1.0}),
                "char_length": ("INT", {"default": 50, "min": 1, "max": 500, "step": 1}),
                "loops": ("INT", {"default": 1, "min": 1, "max": 10, "step": 1}),
                "bright_value": ("INT", {"default": 127, "min": 0, "max": 255, "step": 1}),
                "dark_value": ("INT", {"default": 223, "min": 0, "max": 255, "step": 1}),
                "lower_threshold": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.01}),
                "upper_threshold": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 1.0, "step": 0.01}),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "sort_pixels"
    CATEGORY = "image/effects"

    def sort_pixels(self, image, interval, sorting, angle=0.0, randomness=0.0,
                    char_length=50, loops=1, bright_value=127, dark_value=223,
                    lower_threshold=0.25, upper_threshold=0.8):
        """
        Sort pixels in the image.

        Args:
            image: ComfyUI image tensor (B, H, W, C) in range [0, 1]
            interval: Interval function name
            sorting: Sorting function name
            angle: Rotation angle in degrees
            randomness: Randomness percentage
            char_length: Characteristic length for random/waves
            loops: Number of sorting passes
            bright_value: Threshold for bright mode
            dark_value: Threshold for dark mode
            lower_threshold: Lower threshold for threshold mode
            upper_threshold: Upper threshold for threshold mode

        Returns:
            Sorted image tensor
        """
        # Convert ComfyUI tensor to PIL Image
        # ComfyUI format: (B, H, W, C) in range [0, 1]
        # PIL format: (W, H) in range [0, 255]

        batch_size = image.shape[0]
        results = []

        for i in range(batch_size):
            # Get single image from batch
            img_tensor = image[i]

            # Convert to numpy and scale to [0, 255]
            img_np = (img_tensor.cpu().numpy() * 255).astype(np.uint8)

            # Convert to PIL Image
            pil_img = Image.fromarray(img_np, mode='RGB')

            # Apply pixel sorting
            sorted_pil = pixelsort(
                pil_img,
                interval=interval,
                sorting=sorting,
                angle=angle,
                randomness=randomness,
                char_length=char_length,
                loops=loops,
                bright_value=bright_value,
                dark_value=dark_value,
                lower_threshold=lower_threshold,
                upper_threshold=upper_threshold,
            )

            # Convert back to ComfyUI format
            sorted_np = np.array(sorted_pil.convert('RGB')).astype(np.float32) / 255.0
            sorted_tensor = torch.from_numpy(sorted_np)

            results.append(sorted_tensor)

        # Stack batch
        output = torch.stack(results)

        return (output,)


class PixelSortAnalyzeNode:
    """ComfyUI node to analyze images and suggest thresholds."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
            },
        }

    RETURN_TYPES = ("INT", "INT", "INT", "INT")
    RETURN_NAMES = ("bright_subtle", "bright_moderate", "dark_subtle", "dark_moderate")
    FUNCTION = "analyze_image"
    CATEGORY = "image/analysis"

    def analyze_image(self, image):
        """
        Analyze image and return suggested thresholds.

        Args:
            image: ComfyUI image tensor

        Returns:
            Tuple of suggested threshold values
        """
        # Convert first image in batch to PIL
        img_tensor = image[0]
        img_np = (img_tensor.cpu().numpy() * 255).astype(np.uint8)
        pil_img = Image.fromarray(img_np, mode='RGB')

        # Analyze
        stats = analyze(pil_img)
        suggestions = stats['suggestions']

        return (
            suggestions['bright']['subtle'],
            suggestions['bright']['moderate'],
            suggestions['dark']['subtle'],
            suggestions['dark']['moderate'],
        )


# ComfyUI node registration
NODE_CLASS_MAPPINGS = {
    "PixelSorting": PixelSortNode,
    "PixelSortAnalyze": PixelSortAnalyzeNode,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "PixelSorting": "Pixel Sort Effect",
    "PixelSortAnalyze": "Pixel Sort Analyze",
}