"""
ComfyUI-SmartGridCutter — Universal Grid Image Splitter

Automatically detects grid layout by analyzing contrast/brightness
and splits any grid image into individual frames.

Works with any grid: 3×3 storyboards, 4×2 character sheets,
2×3 panels, or any arbitrary grid layout.

Author: iMind Studio
License: MIT
"""

import torch
import numpy as np
from PIL import Image
import io
import urllib.request


class SmartGridCutter:
    """
    Splits a grid image into individual cells by detecting
    divider lines through brightness/contrast analysis.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "mode": (["auto", "3x3", "4x2", "3x2", "2x2", "custom"],),
            },
            "optional": {
                "custom_cols": ("INT", {"default": 3, "min": 1, "max": 10, "step": 1}),
                "custom_rows": ("INT", {"default": 3, "min": 1, "max": 10, "step": 1}),
                "divider_threshold": ("FLOAT", {"default": 0.15, "min": 0.01, "max": 0.5, "step": 0.01,
                    "tooltip": "Brightness difference threshold for detecting divider lines. Lower = more sensitive."}),
                "min_divider_width": ("INT", {"default": 2, "min": 1, "max": 20, "step": 1,
                    "tooltip": "Minimum width in pixels for a line to be considered a divider."}),
                "padding": ("INT", {"default": 2, "min": 0, "max": 20, "step": 1,
                    "tooltip": "Pixels to trim from each cell edge after cutting."}),
                "align_to": ("INT", {"default": 16, "min": 1, "max": 64, "step": 1,
                    "tooltip": "Align cell dimensions to this value (16 for video compatibility)."}),
            }
        }

    RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE",
                    "IMAGE", "IMAGE", "IMAGE", "IMAGE",
                    "INT", "INT", "INT", "STRING")
    RETURN_NAMES = ("frame_0", "frame_1", "frame_2", "frame_3", "frame_4",
                    "frame_5", "frame_6", "frame_7", "frame_8",
                    "cols", "rows", "total_frames", "grid_info")

    FUNCTION = "cut_grid"
    CATEGORY = "image/transform"
    DESCRIPTION = "Splits a grid image into individual frames by detecting divider lines. Works with any grid layout."

    def _find_dividers(self, gray_array, axis, threshold, min_width, margin_pct=0.05):
        """
        Find divider positions along an axis using variance + contrast analysis.

        Dividers are detected by TWO criteria:
        1. Low variance (uniform color — solid white/black/gray line)
        2. Contrast with neighbors (different brightness from surrounding cells)

        Args:
            gray_array: 2D numpy array (grayscale, 0-1 float)
            axis: 0 for horizontal dividers (scanning rows), 1 for vertical (scanning cols)
            threshold: NOT USED directly anymore, kept for API compat
            min_width: minimum pixel width for a valid divider
            margin_pct: ignore edges (percentage of total size)
        """
        h, w = gray_array.shape

        if axis == 1:  # vertical dividers (scan columns)
            size = w
            # Per-column stats
            profile_mean = np.mean(gray_array[::2, :], axis=0)
            profile_std = np.std(gray_array[::2, :], axis=0)
        else:  # horizontal dividers (scan rows)
            size = h
            profile_mean = np.mean(gray_array[:, ::2], axis=1)
            profile_std = np.std(gray_array[:, ::2], axis=1)

        # Dividers have LOW std (uniform color across the perpendicular axis)
        # Content cells have HIGH std (varied pixel values)
        median_std = np.median(profile_std)
        # A divider line's std should be much lower than typical content
        std_threshold = max(median_std * 0.3, 0.01)

        margin = int(size * margin_pct)
        dividers = []
        in_divider = False
        start = 0

        for i in range(margin, size - margin):
            is_uniform = profile_std[i] < std_threshold

            if is_uniform:
                if not in_divider:
                    start = i
                    in_divider = True
            else:
                if in_divider:
                    width = i - start
                    if width >= min_width:
                        # Verify it's actually a divider by checking contrast with neighbors
                        div_brightness = np.mean(profile_mean[start:i])
                        # Check brightness of cells on both sides
                        left_zone = max(0, start - min(30, start))
                        right_zone = min(size, i + 30)
                        neighbor_brightness = np.mean(
                            np.concatenate([
                                profile_mean[left_zone:start],
                                profile_mean[i:right_zone]
                            ])
                        ) if (start > left_zone or i < right_zone) else div_brightness

                        contrast = abs(div_brightness - neighbor_brightness)
                        if contrast > 0.05:  # Must have SOME contrast with neighbors
                            dividers.append({
                                "start": start,
                                "end": i,
                                "mid": (start + i) // 2,
                                "width": width,
                            })
                    in_divider = False

        return dividers

    def _build_grid_from_dividers(self, x_dividers, y_dividers, img_w, img_h):
        """Build cell boundaries from detected dividers."""
        cols = len(x_dividers) + 1
        rows = len(y_dividers) + 1

        # X boundaries: [0, end_of_div1, end_of_div2, ..., img_w]
        x_starts = [0] + [d["end"] for d in x_dividers]
        x_ends = [d["start"] for d in x_dividers] + [img_w]

        y_starts = [0] + [d["end"] for d in y_dividers]
        y_ends = [d["start"] for d in y_dividers] + [img_h]

        return cols, rows, x_starts, x_ends, y_starts, y_ends

    def _build_grid_from_preset(self, cols, rows, img_w, img_h):
        """Build cell boundaries using simple division."""
        cell_w = img_w // cols
        cell_h = img_h // rows

        x_starts = [i * cell_w for i in range(cols)]
        x_ends = [(i + 1) * cell_w for i in range(cols)]
        y_starts = [i * cell_h for i in range(rows)]
        y_ends = [(i + 1) * cell_h for i in range(rows)]

        return x_starts, x_ends, y_starts, y_ends

    def _tensor_to_pil(self, tensor):
        """Convert ComfyUI IMAGE tensor to PIL Image."""
        # ComfyUI images: [B, H, W, C], float32, 0-1
        if tensor.dim() == 4:
            tensor = tensor[0]  # Take first batch
        arr = (tensor.cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
        return Image.fromarray(arr)

    def _pil_to_tensor(self, pil_img):
        """Convert PIL Image to ComfyUI IMAGE tensor."""
        arr = np.array(pil_img).astype(np.float32) / 255.0
        if arr.ndim == 2:
            arr = np.stack([arr] * 3, axis=-1)
        return torch.from_numpy(arr).unsqueeze(0)  # [1, H, W, C]

    def cut_grid(self, image, mode, custom_cols=3, custom_rows=3,
                 divider_threshold=0.15, min_divider_width=2,
                 padding=2, align_to=16):
        """Main cutting function."""

        pil_img = self._tensor_to_pil(image)
        img_w, img_h = pil_img.size

        # Convert to grayscale for analysis
        gray = np.array(pil_img.convert("L")).astype(np.float32) / 255.0

        # Determine grid layout
        presets = {
            "3x3": (3, 3),
            "4x2": (4, 2),
            "3x2": (3, 2),
            "2x2": (2, 2),
            "custom": (custom_cols, custom_rows),
        }

        if mode == "auto":
            # Smart detection
            x_dividers = self._find_dividers(gray, axis=1,
                threshold=divider_threshold, min_width=min_divider_width)
            y_dividers = self._find_dividers(gray, axis=0,
                threshold=divider_threshold, min_width=min_divider_width)

            cols, rows, x_starts, x_ends, y_starts, y_ends = \
                self._build_grid_from_dividers(x_dividers, y_dividers, img_w, img_h)

            detect_method = "auto_detected"

            # Sanity check: if detection gives weird results (1×1 or 20×20), fallback
            if cols < 2 and rows < 2:
                # Try preset-based detection using aspect ratio
                ar = img_w / img_h
                if 0.9 <= ar <= 1.1:
                    cols, rows = 3, 3  # Square → likely 3×3
                elif ar > 1.4:
                    cols, rows = 4, 2  # Wide → likely 4×2
                else:
                    cols, rows = 3, 2  # Portrait-ish → likely 3×2

                x_starts, x_ends, y_starts, y_ends = \
                    self._build_grid_from_preset(cols, rows, img_w, img_h)
                detect_method = "aspect_ratio_fallback"

        else:
            cols, rows = presets[mode]
            # Try smart detection first, validate against preset
            x_dividers = self._find_dividers(gray, axis=1,
                threshold=divider_threshold, min_width=min_divider_width)
            y_dividers = self._find_dividers(gray, axis=0,
                threshold=divider_threshold, min_width=min_divider_width)

            detected_cols = len(x_dividers) + 1
            detected_rows = len(y_dividers) + 1

            if detected_cols == cols and detected_rows == rows:
                # Smart detection matches preset — use smart bounds
                _, _, x_starts, x_ends, y_starts, y_ends = \
                    self._build_grid_from_dividers(x_dividers, y_dividers, img_w, img_h)
                detect_method = "smart_verified"
            else:
                # Fallback to simple division
                x_starts, x_ends, y_starts, y_ends = \
                    self._build_grid_from_preset(cols, rows, img_w, img_h)
                detect_method = f"preset_fallback (detected {detected_cols}x{detected_rows})"

        # Cut frames
        frames = []
        for r in range(rows):
            for c in range(cols):
                left = x_starts[c] + padding
                top = y_starts[r] + padding
                right = x_ends[c] - padding
                bottom = y_ends[r] - padding

                w = right - left
                h = bottom - top

                # Align dimensions
                if align_to > 1:
                    w = (w // align_to) * align_to
                    h = (h // align_to) * align_to

                cell = pil_img.crop((left, top, left + w, top + h))
                frames.append(self._pil_to_tensor(cell))

        total = len(frames)

        # Build info string
        info = (f"Grid: {cols}×{rows} = {total} frames | "
                f"Source: {img_w}×{img_h} | "
                f"Cell: {frames[0].shape[2]}×{frames[0].shape[1]} | "
                f"Detection: {detect_method}")

        # Pad to 9 outputs (max supported)
        while len(frames) < 9:
            # Create a 1×1 black tensor as placeholder
            frames.append(torch.zeros(1, 16, 16, 3))

        return (*frames[:9], cols, rows, total, info)


class SmartGridCutterURL:
    """
    Downloads an image from URL and splits it into grid cells.
    Useful for n8n / API integration.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image_url": ("STRING", {"default": "", "multiline": False}),
                "mode": (["auto", "3x3", "4x2", "3x2", "2x2", "custom"],),
            },
            "optional": {
                "custom_cols": ("INT", {"default": 3, "min": 1, "max": 10}),
                "custom_rows": ("INT", {"default": 3, "min": 1, "max": 10}),
                "divider_threshold": ("FLOAT", {"default": 0.15, "min": 0.01, "max": 0.5, "step": 0.01}),
                "padding": ("INT", {"default": 2, "min": 0, "max": 20}),
                "align_to": ("INT", {"default": 16, "min": 1, "max": 64}),
            }
        }

    RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE",
                    "IMAGE", "IMAGE", "IMAGE", "IMAGE",
                    "INT", "INT", "INT", "STRING")
    RETURN_NAMES = ("frame_0", "frame_1", "frame_2", "frame_3", "frame_4",
                    "frame_5", "frame_6", "frame_7", "frame_8",
                    "cols", "rows", "total_frames", "grid_info")

    FUNCTION = "cut_from_url"
    CATEGORY = "image/transform"
    DESCRIPTION = "Downloads image from URL and splits it into grid frames."

    def cut_from_url(self, image_url, mode, custom_cols=3, custom_rows=3,
                     divider_threshold=0.15, padding=2, align_to=16):
        """Download image from URL and cut it."""
        # Download image
        req = urllib.request.Request(image_url, headers={"User-Agent": "ComfyUI-SmartGridCutter/1.0"})
        with urllib.request.urlopen(req, timeout=30) as response:
            img_data = response.read()

        pil_img = Image.open(io.BytesIO(img_data)).convert("RGB")

        # Convert to tensor
        arr = np.array(pil_img).astype(np.float32) / 255.0
        tensor = torch.from_numpy(arr).unsqueeze(0)

        # Use the main cutter
        cutter = SmartGridCutter()
        return cutter.cut_grid(tensor, mode, custom_cols, custom_rows,
                              divider_threshold, 2, padding, align_to)


# ── Registration ─────────────────────────────────────
NODE_CLASS_MAPPINGS = {
    "SmartGridCutter": SmartGridCutter,
    "SmartGridCutterURL": SmartGridCutterURL,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "SmartGridCutter": "Smart Grid Cutter ✂️",
    "SmartGridCutterURL": "Smart Grid Cutter (URL) ✂️🔗",
}