import json
import math
from typing import List, Sequence, Tuple

import torch
import torch.nn.functional as F


CATEGORY = "Universal Remove"


def _ensure_image_batch(images: torch.Tensor) -> torch.Tensor:
    if not isinstance(images, torch.Tensor):
        raise TypeError("Expected IMAGE input as torch.Tensor")
    if images.ndim == 3:
        images = images.unsqueeze(0)
    if images.ndim != 4:
        raise ValueError(f"Expected IMAGE tensor with 4 dims, got {images.ndim}")
    return images.clamp(0.0, 1.0)


def _ensure_mask_batch(mask: torch.Tensor) -> torch.Tensor:
    if not isinstance(mask, torch.Tensor):
        raise TypeError("Expected MASK input as torch.Tensor")
    if mask.ndim == 2:
        mask = mask.unsqueeze(0)
    elif mask.ndim == 4 and mask.shape[1] == 1:
        mask = mask[:, 0]
    elif mask.ndim == 4 and mask.shape[-1] == 1:
        mask = mask[..., 0]
    if mask.ndim != 3:
        raise ValueError(f"Expected MASK tensor with 3 dims, got {mask.ndim}")
    return mask.float().clamp(0.0, 1.0)


def _tile_to_batch(tensor: torch.Tensor, target_batch: int) -> torch.Tensor:
    if tensor.shape[0] == target_batch:
        return tensor
    if tensor.shape[0] <= 0:
        raise ValueError("Cannot tile an empty batch")
    repeats = math.ceil(target_batch / tensor.shape[0])
    repeat_shape = [repeats] + [1] * (tensor.ndim - 1)
    return tensor.repeat(*repeat_shape)[:target_batch]


def _match_mask_to_images(mask: torch.Tensor, images: torch.Tensor) -> torch.Tensor:
    mask = _ensure_mask_batch(mask)
    images = _ensure_image_batch(images)
    if mask.shape[0] != images.shape[0]:
        mask = _tile_to_batch(mask, images.shape[0])
    if mask.shape[1] != images.shape[1] or mask.shape[2] != images.shape[2]:
        mask = F.interpolate(
            mask.unsqueeze(1),
            size=(images.shape[1], images.shape[2]),
            mode="nearest",
        ).squeeze(1)
    return mask.clamp(0.0, 1.0)


def _to_bchw(images: torch.Tensor) -> torch.Tensor:
    return images.permute(0, 3, 1, 2)


def _to_bhwc(images: torch.Tensor) -> torch.Tensor:
    return images.permute(0, 2, 3, 1)


def _resize_images(images: torch.Tensor, size: Tuple[int, int]) -> torch.Tensor:
    bchw = _to_bchw(images)
    resized = F.interpolate(bchw, size=size, mode="bilinear", align_corners=False)
    return _to_bhwc(resized).clamp(0.0, 1.0)


def _resize_masks(mask: torch.Tensor, size: Tuple[int, int]) -> torch.Tensor:
    resized = F.interpolate(mask.unsqueeze(1), size=size, mode="nearest")
    return resized.squeeze(1).clamp(0.0, 1.0)


def _pad_images(images: torch.Tensor, target_h: int, target_w: int) -> torch.Tensor:
    pad_h = max(0, target_h - images.shape[1])
    pad_w = max(0, target_w - images.shape[2])
    if pad_h == 0 and pad_w == 0:
        return images
    bchw = _to_bchw(images)
    padded = F.pad(bchw, (0, pad_w, 0, pad_h), mode="constant", value=0.0)
    return _to_bhwc(padded).clamp(0.0, 1.0)


def _pad_masks(mask: torch.Tensor, target_h: int, target_w: int) -> torch.Tensor:
    pad_h = max(0, target_h - mask.shape[1])
    pad_w = max(0, target_w - mask.shape[2])
    if pad_h == 0 and pad_w == 0:
        return mask
    padded = F.pad(mask.unsqueeze(1), (0, pad_w, 0, pad_h), mode="constant", value=0.0)
    return padded.squeeze(1).clamp(0.0, 1.0)


def _parse_hex_color(color: str) -> torch.Tensor:
    text = color.strip()
    if text.startswith("#"):
        text = text[1:]
    if len(text) != 6:
        raise ValueError("Color must be a 6-digit hex string, e.g. #00FF00")
    rgb = [int(text[i : i + 2], 16) / 255.0 for i in range(0, 6, 2)]
    return torch.tensor(rgb, dtype=torch.float32)


def _overlay_mask(images: torch.Tensor, mask: torch.Tensor, color: torch.Tensor, opacity: float) -> torch.Tensor:
    images = _ensure_image_batch(images)
    mask = _match_mask_to_images(mask, images)
    color = color.to(device=images.device, dtype=images.dtype).view(1, 1, 1, 3)
    alpha = mask.unsqueeze(-1).clamp(0.0, 1.0) * float(opacity)
    return (images * (1.0 - alpha) + color * alpha).clamp(0.0, 1.0)


def _gaussian_kernel_1d(kernel_size: int, sigma: float, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
    if kernel_size % 2 == 0:
        kernel_size += 1
    radius = kernel_size // 2
    if sigma <= 0:
        sigma = max(float(kernel_size) / 6.0, 1e-6)
    coords = torch.arange(-radius, radius + 1, device=device, dtype=dtype)
    kernel = torch.exp(-0.5 * (coords / sigma) ** 2)
    kernel = kernel / kernel.sum()
    return kernel


def _gaussian_blur_mask(mask: torch.Tensor, kernel_size: int, sigma: float) -> torch.Tensor:
    mask = _ensure_mask_batch(mask)
    if kernel_size <= 1 and sigma <= 0:
        return mask
    if kernel_size <= 1:
        kernel_size = int(max(3, math.ceil(sigma * 6)))
    if kernel_size % 2 == 0:
        kernel_size += 1

    kernel_1d = _gaussian_kernel_1d(kernel_size, sigma, mask.device, mask.dtype)
    kernel_x = kernel_1d.view(1, 1, 1, -1)
    kernel_y = kernel_1d.view(1, 1, -1, 1)

    blurred = mask.unsqueeze(1)
    pad = kernel_size // 2
    blurred = F.pad(blurred, (pad, pad, 0, 0), mode="replicate")
    blurred = F.conv2d(blurred, kernel_x)
    blurred = F.pad(blurred, (0, 0, pad, pad), mode="replicate")
    blurred = F.conv2d(blurred, kernel_y)
    return blurred.squeeze(1).clamp(0.0, 1.0)


def _expand_mask(mask: torch.Tensor, expand_pixels: int) -> torch.Tensor:
    mask = _ensure_mask_batch(mask)
    if expand_pixels <= 0:
        return mask
    kernel = expand_pixels * 2 + 1
    expanded = F.max_pool2d(mask.unsqueeze(1), kernel_size=kernel, stride=1, padding=expand_pixels)
    return expanded.squeeze(1).clamp(0.0, 1.0)


def _fill_masked_region(images: torch.Tensor, mask: torch.Tensor, color: torch.Tensor) -> torch.Tensor:
    images = _ensure_image_batch(images)
    mask = _match_mask_to_images(mask, images)
    color = color.to(device=images.device, dtype=images.dtype).view(1, 1, 1, 3)
    return (images * (1.0 - mask.unsqueeze(-1)) + color * mask.unsqueeze(-1)).clamp(0.0, 1.0)


def _normalize_box_values(values: Sequence[float], width: int, height: int, normalized: bool) -> Tuple[float, float, float, float]:
    x1, y1, x2, y2 = values
    if normalized:
        x1 *= width
        x2 *= width
        y1 *= height
        y2 *= height
    return x1, y1, x2, y2


def _parse_single_box(item, coordinate_format: str, width: int, height: int, normalized: bool) -> Tuple[int, int, int, int]:
    if isinstance(item, dict):
        if {"x", "y", "w", "h"}.issubset(item):
            x1 = float(item["x"])
            y1 = float(item["y"])
            w = float(item["w"])
            h = float(item["h"])
            if normalized:
                x1 *= width
                y1 *= height
                w *= width
                h *= height
            x2 = x1 + w
            y2 = y1 + h
        elif {"x1", "y1", "x2", "y2"}.issubset(item):
            x1, y1, x2, y2 = _normalize_box_values(
                [float(item["x1"]), float(item["y1"]), float(item["x2"]), float(item["y2"])],
                width,
                height,
                normalized,
            )
        elif {"left", "top", "right", "bottom"}.issubset(item):
            x1, y1, x2, y2 = _normalize_box_values(
                [float(item["left"]), float(item["top"]), float(item["right"]), float(item["bottom"])],
                width,
                height,
                normalized,
            )
        else:
            raise ValueError(f"Unsupported box object: {item}")
    elif isinstance(item, (list, tuple)) and len(item) == 4:
        values = [float(v) for v in item]
        if coordinate_format == "xywh":
            x1, y1, w, h = values
            if normalized:
                x1 *= width
                y1 *= height
                w *= width
                h *= height
            x2 = x1 + w
            y2 = y1 + h
        else:
            x1, y1, x2, y2 = _normalize_box_values(values, width, height, normalized)
    else:
        raise ValueError(f"Unsupported box entry: {item}")

    left = int(math.floor(min(x1, x2)))
    right = int(math.ceil(max(x1, x2)))
    top = int(math.floor(min(y1, y2)))
    bottom = int(math.ceil(max(y1, y2)))

    left = max(0, min(width, left))
    right = max(0, min(width, right))
    top = max(0, min(height, top))
    bottom = max(0, min(height, bottom))
    return left, top, right, bottom


def _parse_boxes_json(boxes_json: str, coordinate_format: str, width: int, height: int, normalized: bool) -> List[Tuple[int, int, int, int]]:
    text = boxes_json.strip()
    if not text:
        return []
    data = json.loads(text)
    if isinstance(data, dict):
        if "boxes" in data:
            data = data["boxes"]
        else:
            data = [data]
    if not isinstance(data, list):
        raise ValueError("boxes_json must be a JSON array or object")
    boxes = []
    for item in data:
        box = _parse_single_box(item, coordinate_format, width, height, normalized)
        if box[2] > box[0] and box[3] > box[1]:
            boxes.append(box)
    return boxes


def _scale_point(x: float, y: float, width: int, height: int, normalized: bool) -> Tuple[float, float]:
    if normalized:
        x *= width
        y *= height
    return x, y


def _region_frame_index(region: dict, default_frame_index: int, batch_size: int) -> int:
    index = int(region.get("frame_index", default_frame_index))
    return max(0, min(index, batch_size - 1))


def _region_type(region: dict) -> str:
    value = str(region.get("type", region.get("shape", "rect"))).strip().lower()
    aliases = {
        "rectangle": "rect",
        "box": "rect",
        "bbox": "rect",
        "ellipse": "circle",
        "poly": "polygon",
    }
    return aliases.get(value, value)


def _polygon_to_mask(points: Sequence[Sequence[float]], width: int, height: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
    if len(points) < 3:
        return torch.zeros((height, width), dtype=dtype, device=device)

    px = torch.arange(width, dtype=dtype, device=device).view(1, width) + 0.5
    py = torch.arange(height, dtype=dtype, device=device).view(height, 1) + 0.5
    inside = torch.zeros((height, width), dtype=torch.bool, device=device)

    wrapped = list(points) + [points[0]]
    eps = torch.tensor(1e-6, dtype=dtype, device=device)
    for current, nxt in zip(wrapped[:-1], wrapped[1:]):
        x1 = torch.tensor(float(current[0]), dtype=dtype, device=device)
        y1 = torch.tensor(float(current[1]), dtype=dtype, device=device)
        x2 = torch.tensor(float(nxt[0]), dtype=dtype, device=device)
        y2 = torch.tensor(float(nxt[1]), dtype=dtype, device=device)
        crosses_y = torch.logical_xor(y1 > py, y2 > py)
        x_intersection = (x2 - x1) * (py - y1) / (y2 - y1 + eps) + x1
        intersects = crosses_y & (px < x_intersection)
        inside = torch.logical_xor(inside, intersects)

    return inside.to(dtype=dtype)


def _draw_region(mask: torch.Tensor, region: dict, width: int, height: int, normalized: bool) -> None:
    shape = _region_type(region)
    if shape == "rect":
        box = _parse_single_box(region, "xyxy", width, height, normalized)
        x1, y1, x2, y2 = box
        if x2 > x1 and y2 > y1:
            mask[y1:y2, x1:x2] = 1.0
        return

    if shape == "circle":
        if "center" in region:
            cx, cy = region["center"]
        else:
            cx = region.get("cx", region.get("x", 0.0))
            cy = region.get("cy", region.get("y", 0.0))
        cx, cy = _scale_point(float(cx), float(cy), width, height, normalized)
        radius = float(region.get("r", region.get("radius", 0.0)))
        if normalized:
            radius *= min(width, height)
        if radius <= 0:
            return
        yy, xx = torch.meshgrid(
            torch.arange(height, dtype=mask.dtype, device=mask.device),
            torch.arange(width, dtype=mask.dtype, device=mask.device),
            indexing="ij",
        )
        circle = (xx + 0.5 - cx) ** 2 + (yy + 0.5 - cy) ** 2 <= radius ** 2
        mask[circle] = 1.0
        return

    if shape == "polygon":
        raw_points = region.get("points", [])
        points = []
        for point in raw_points:
            if isinstance(point, dict):
                x = float(point.get("x", 0.0))
                y = float(point.get("y", 0.0))
            else:
                x = float(point[0])
                y = float(point[1])
            points.append(_scale_point(x, y, width, height, normalized))
        mask.copy_(torch.maximum(mask, _polygon_to_mask(points, width, height, mask.device, mask.dtype)))
        return

    raise ValueError(f"Unsupported region type: {shape}")


def _parse_regions_json(regions_json: str) -> List[dict]:
    text = regions_json.strip()
    if not text:
        return []
    data = json.loads(text)
    if isinstance(data, dict):
        if "regions" in data:
            data = data["regions"]
        else:
            data = [data]
    if not isinstance(data, list):
        raise ValueError("regions_json must be a JSON array or object")
    for region in data:
        if not isinstance(region, dict):
            raise ValueError(f"Region entries must be objects, got: {region}")
    return data


def _apply_keyframe_masks(keyframe_masks: dict, batch: int, height: int, width: int, dtype: torch.dtype, device: torch.device, broadcast_mode: str) -> torch.Tensor:
    mask = torch.zeros((batch, height, width), dtype=dtype, device=device)
    if not keyframe_masks:
        return mask

    sorted_items = sorted(keyframe_masks.items())
    if broadcast_mode == "repeat_to_batch":
        combined = torch.zeros((height, width), dtype=dtype, device=device)
        for _, key_mask in sorted_items:
            combined = torch.maximum(combined, key_mask)
        return combined.unsqueeze(0).repeat(batch, 1, 1)

    if broadcast_mode == "hold_until_next_keyframe":
        for idx, (start_frame, key_mask) in enumerate(sorted_items):
            end_frame = sorted_items[idx + 1][0] if idx + 1 < len(sorted_items) else batch
            if end_frame > start_frame:
                mask[start_frame:end_frame] = key_mask
        return mask

    for frame, key_mask in sorted_items:
        if broadcast_mode == "hold_after_frame":
            mask[frame:] = key_mask
        else:
            mask[frame] = key_mask
    return mask


class URDemoImage:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "width": ("INT", {"default": 768, "min": 64, "max": 8192, "step": 8}),
                "height": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 8}),
                "batch_size": ("INT", {"default": 1, "min": 1, "max": 256}),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "generate"
    CATEGORY = CATEGORY

    def generate(self, width: int, height: int, batch_size: int):
        x = torch.linspace(0.0, 1.0, int(width), dtype=torch.float32).view(1, 1, width)
        y = torch.linspace(0.0, 1.0, int(height), dtype=torch.float32).view(1, height, 1)
        red = 0.15 + 0.35 * x.expand(1, height, width)
        green = 0.18 + 0.45 * y.expand(1, height, width)
        blue = 0.25 + 0.25 * (1.0 - x.expand(1, height, width))
        image = torch.stack([red, green, blue], dim=-1).repeat(int(batch_size), 1, 1, 1)

        x1 = max(0, min(int(width), int(width * 0.22)))
        x2 = max(0, min(int(width), int(width * 0.55)))
        y1 = max(0, min(int(height), int(height * 0.24)))
        y2 = max(0, min(int(height), int(height * 0.50)))
        image[:, y1:y2, x1:x2, 0] = 0.92
        image[:, y1:y2, x1:x2, 1] = 0.16
        image[:, y1:y2, x1:x2, 2] = 0.12
        return (image.clamp(0.0, 1.0),)


class URRemovePromptBuilder:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "target_prompt": ("STRING", {"multiline": False, "default": "object to remove"}),
                "scene_hint": ("STRING", {"multiline": True, "default": ""}),
                "preserve_prompt": ("STRING", {"multiline": True, "default": ""}),
                "extra_negative_prompt": ("STRING", {"multiline": True, "default": ""}),
                "fill_mode": (["background_reconstruct", "clean_plate", "direct_replace"],),
                "masked_edit_only": ("BOOLEAN", {"default": True}),
            }
        }

    RETURN_TYPES = ("STRING", "STRING", "STRING")
    RETURN_NAMES = ("segmentation_prompt", "remove_positive_prompt", "remove_negative_prompt")
    FUNCTION = "build"
    CATEGORY = CATEGORY

    def build(
        self,
        target_prompt: str,
        scene_hint: str,
        preserve_prompt: str,
        extra_negative_prompt: str,
        fill_mode: str,
        masked_edit_only: bool,
    ):
        target = target_prompt.strip() or "target object"
        segmentation_prompt = target

        positive_parts = [f"Remove {target} from the masked region."]
        if fill_mode == "background_reconstruct":
            positive_parts.append("Reconstruct a plausible background behind the removed content.")
        elif fill_mode == "clean_plate":
            positive_parts.append("Create a clean plate result that matches the surrounding scene.")
        else:
            positive_parts.append("Replace the masked region with visually coherent scene content.")

        if masked_edit_only:
            positive_parts.append("Only modify the masked region and keep unmasked content unchanged.")
        if scene_hint.strip():
            positive_parts.append(f"Scene context: {scene_hint.strip()}.")
        if preserve_prompt.strip():
            positive_parts.append(f"Preserve: {preserve_prompt.strip()}.")

        negative_parts = [
            "artifacts",
            "ghosting",
            "duplicate subject",
            "warped anatomy",
            "broken perspective",
            "extra text",
            "watermark",
            "subtitle",
            "logo",
            "blur",
            "low quality",
        ]
        if extra_negative_prompt.strip():
            negative_parts.append(extra_negative_prompt.strip())

        return (
            segmentation_prompt,
            " ".join(positive_parts),
            ", ".join(negative_parts),
        )


class URRegionEditor:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "regions_json": (
                    "STRING",
                    {
                        "multiline": True,
                        "default": "[]",
                    },
                ),
                "canvas_width": ("INT", {"default": 768, "min": 64, "max": 8192, "step": 8}),
                "canvas_height": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 8}),
                "default_frame_index": ("INT", {"default": 0, "min": 0, "max": 1000000}),
            }
        }

    RETURN_TYPES = ("STRING", "INT", "INT", "INT")
    RETURN_NAMES = ("regions_json", "canvas_width", "canvas_height", "region_count")
    FUNCTION = "emit"
    CATEGORY = CATEGORY

    def emit(self, regions_json: str, canvas_width: int, canvas_height: int, default_frame_index: int):
        regions = _parse_regions_json(regions_json)
        normalized_regions = []
        for region in regions:
            copied = dict(region)
            copied["frame_index"] = int(copied.get("frame_index", default_frame_index))
            normalized_regions.append(copied)
        text = json.dumps(normalized_regions, ensure_ascii=False)
        return (text, int(canvas_width), int(canvas_height), len(normalized_regions))


class URBoxesToMaskBatch:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE",),
                "boxes_json": ("STRING", {"multiline": True, "default": "[]"}),
                "frame_index": ("INT", {"default": 0, "min": 0, "max": 1000000}),
                "coordinate_format": (["xyxy", "xywh"],),
                "broadcast_mode": (["repeat_to_batch", "single_frame_only", "hold_after_frame"],),
                "normalized": ("BOOLEAN", {"default": False}),
                "preview_color": ("STRING", {"multiline": False, "default": "#FF3B30"}),
                "preview_opacity": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01}),
            }
        }

    RETURN_TYPES = ("MASK", "IMAGE", "IMAGE")
    RETURN_NAMES = ("mask", "selected_frame", "preview")
    FUNCTION = "build_masks"
    CATEGORY = CATEGORY

    def build_masks(
        self,
        images: torch.Tensor,
        boxes_json: str,
        frame_index: int,
        coordinate_format: str,
        broadcast_mode: str,
        normalized: bool,
        preview_color: str,
        preview_opacity: float,
    ):
        images = _ensure_image_batch(images)
        batch, height, width, _ = images.shape
        index = max(0, min(int(frame_index), batch - 1))

        boxes = _parse_boxes_json(boxes_json, coordinate_format, width, height, normalized)
        selected_mask = torch.zeros((1, height, width), dtype=images.dtype, device=images.device)
        for x1, y1, x2, y2 in boxes:
            selected_mask[:, y1:y2, x1:x2] = 1.0

        if broadcast_mode == "repeat_to_batch":
            mask = selected_mask.repeat(batch, 1, 1)
        elif broadcast_mode == "hold_after_frame":
            mask = torch.zeros((batch, height, width), dtype=images.dtype, device=images.device)
            mask[index:] = selected_mask
        else:
            mask = torch.zeros((batch, height, width), dtype=images.dtype, device=images.device)
            mask[index] = selected_mask[0]

        preview = _overlay_mask(
            images[index : index + 1],
            selected_mask,
            _parse_hex_color(preview_color),
            preview_opacity,
        )
        return (mask, images[index : index + 1], preview)


class URRegionsToMaskBatch:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE",),
                "regions_json": (
                    "STRING",
                    {
                        "multiline": True,
                        "default": '[{"type":"rect","frame_index":0,"x":160,"y":120,"w":260,"h":140}]',
                    },
                ),
                "default_frame_index": ("INT", {"default": 0, "min": 0, "max": 1000000}),
                "broadcast_mode": (["repeat_to_batch", "single_frame_only", "hold_after_frame", "hold_until_next_keyframe"],),
                "normalized": ("BOOLEAN", {"default": False}),
                "preview_frame": ("INT", {"default": 0, "min": 0, "max": 1000000}),
                "preview_color": ("STRING", {"multiline": False, "default": "#FF3B30"}),
                "preview_opacity": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01}),
            }
        }

    RETURN_TYPES = ("MASK", "IMAGE", "IMAGE")
    RETURN_NAMES = ("mask", "selected_frame", "preview")
    FUNCTION = "build_masks"
    CATEGORY = CATEGORY

    def build_masks(
        self,
        images: torch.Tensor,
        regions_json: str,
        default_frame_index: int,
        broadcast_mode: str,
        normalized: bool,
        preview_frame: int,
        preview_color: str,
        preview_opacity: float,
    ):
        images = _ensure_image_batch(images)
        batch, height, width, _ = images.shape
        default_index = max(0, min(int(default_frame_index), batch - 1))
        preview_index = max(0, min(int(preview_frame), batch - 1))

        keyframe_masks = {}
        for region in _parse_regions_json(regions_json):
            frame = _region_frame_index(region, default_index, batch)
            if frame not in keyframe_masks:
                keyframe_masks[frame] = torch.zeros((height, width), dtype=images.dtype, device=images.device)
            _draw_region(keyframe_masks[frame], region, width, height, normalized)

        mask = _apply_keyframe_masks(
            keyframe_masks,
            batch,
            height,
            width,
            images.dtype,
            images.device,
            broadcast_mode,
        )
        preview = _overlay_mask(
            images[preview_index : preview_index + 1],
            mask[preview_index : preview_index + 1],
            _parse_hex_color(preview_color),
            preview_opacity,
        )
        return (mask, images[preview_index : preview_index + 1], preview)


class URMaskProcessor:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "mask": ("MASK",),
                "expand_pixels": ("INT", {"default": 0, "min": 0, "max": 2048}),
                "blur_sigma": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 512.0, "step": 0.01}),
                "blur_kernel": ("INT", {"default": 0, "min": 0, "max": 2049}),
                "preview_color": ("STRING", {"multiline": False, "default": "#00FF00"}),
                "preview_opacity": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01}),
            },
            "optional": {
                "image": ("IMAGE",),
            },
        }

    RETURN_TYPES = ("MASK", "IMAGE", "IMAGE")
    RETURN_NAMES = ("mask", "mask_image", "preview")
    FUNCTION = "process"
    CATEGORY = CATEGORY

    def process(
        self,
        mask: torch.Tensor,
        expand_pixels: int,
        blur_sigma: float,
        blur_kernel: int,
        preview_color: str,
        preview_opacity: float,
        image: torch.Tensor = None,
    ):
        mask = _ensure_mask_batch(mask)
        if expand_pixels > 0:
            mask = _expand_mask(mask, expand_pixels)
        if blur_sigma > 0.0 or blur_kernel > 1:
            mask = _gaussian_blur_mask(mask, blur_kernel, blur_sigma)
        mask = mask.clamp(0.0, 1.0)

        mask_image = mask.unsqueeze(-1).repeat(1, 1, 1, 3)
        if image is not None:
            image = _ensure_image_batch(image)
            preview = _overlay_mask(image, mask, _parse_hex_color(preview_color), preview_opacity)
        else:
            preview = _overlay_mask(mask_image, mask, _parse_hex_color(preview_color), preview_opacity)
        return (mask, mask_image, preview)


class URKeyframePrep:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE",),
                "mask": ("MASK",),
                "frame_index": ("INT", {"default": 0, "min": 0, "max": 1000000}),
                "fill_color": ("STRING", {"multiline": False, "default": "#00FF00"}),
                "resize_longest_side": ("INT", {"default": 0, "min": 0, "max": 8192, "step": 8}),
                "pad_to_multiple": ("INT", {"default": 16, "min": 1, "max": 256}),
            }
        }

    RETURN_TYPES = ("IMAGE", "IMAGE", "MASK", "IMAGE", "INT", "INT", "INT")
    RETURN_NAMES = (
        "original_frame",
        "prepared_frame",
        "prepared_mask",
        "filled_frame",
        "prepared_width",
        "prepared_height",
        "frame_count",
    )
    FUNCTION = "prepare"
    CATEGORY = CATEGORY

    def prepare(
        self,
        images: torch.Tensor,
        mask: torch.Tensor,
        frame_index: int,
        fill_color: str,
        resize_longest_side: int,
        pad_to_multiple: int,
    ):
        images = _ensure_image_batch(images)
        mask = _match_mask_to_images(mask, images)
        batch, height, width, _ = images.shape
        index = max(0, min(int(frame_index), batch - 1))

        original_frame = images[index : index + 1]
        prepared_frame = original_frame.clone()
        prepared_mask = mask[index : index + 1].clone()

        if resize_longest_side > 0:
            longest = max(height, width)
            scale = float(resize_longest_side) / float(longest)
            target_h = max(1, int(round(height * scale)))
            target_w = max(1, int(round(width * scale)))
            prepared_frame = _resize_images(prepared_frame, (target_h, target_w))
            prepared_mask = _resize_masks(prepared_mask, (target_h, target_w))

        prepared_h = prepared_frame.shape[1]
        prepared_w = prepared_frame.shape[2]
        if pad_to_multiple > 1:
            target_h = int(math.ceil(prepared_h / pad_to_multiple) * pad_to_multiple)
            target_w = int(math.ceil(prepared_w / pad_to_multiple) * pad_to_multiple)
            prepared_frame = _pad_images(prepared_frame, target_h, target_w)
            prepared_mask = _pad_masks(prepared_mask, target_h, target_w)

        color = _parse_hex_color(fill_color)
        filled_frame = _fill_masked_region(prepared_frame, prepared_mask, color)
        return (
            original_frame,
            prepared_frame,
            prepared_mask,
            filled_frame,
            prepared_frame.shape[2],
            prepared_frame.shape[1],
            batch,
        )


class URSoftComposite:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "inpaint": ("IMAGE",),
                "original": ("IMAGE",),
                "mask": ("MASK",),
                "kernel": ("INT", {"default": 15, "min": 1, "max": 2049}),
                "sigma": ("FLOAT", {"default": 7.5, "min": 0.0, "max": 512.0, "step": 0.01}),
            }
        }

    RETURN_TYPES = ("IMAGE", "MASK")
    RETURN_NAMES = ("image", "soft_mask")
    FUNCTION = "composite"
    CATEGORY = CATEGORY

    def composite(
        self,
        inpaint: torch.Tensor,
        original: torch.Tensor,
        mask: torch.Tensor,
        kernel: int,
        sigma: float,
    ):
        original = _ensure_image_batch(original)
        inpaint = _ensure_image_batch(inpaint)
        target_batch = max(original.shape[0], inpaint.shape[0], _ensure_mask_batch(mask).shape[0])

        original = _tile_to_batch(original, target_batch)
        inpaint = _tile_to_batch(inpaint, target_batch)
        mask = _tile_to_batch(_ensure_mask_batch(mask), target_batch)
        mask = _match_mask_to_images(mask, original)

        if inpaint.shape[1] != original.shape[1] or inpaint.shape[2] != original.shape[2]:
            inpaint = _resize_images(inpaint, (original.shape[1], original.shape[2]))

        soft_mask = _gaussian_blur_mask(mask, kernel, sigma).clamp(0.0, 1.0)
        composited = original * (1.0 - soft_mask.unsqueeze(-1)) + inpaint * soft_mask.unsqueeze(-1)
        return (composited.clamp(0.0, 1.0), soft_mask)


NODE_CLASS_MAPPINGS = {
    "UR Demo Image": URDemoImage,
    "UR Remove Prompt Builder": URRemovePromptBuilder,
    "UR Region Editor": URRegionEditor,
    "UR Boxes To Mask Batch": URBoxesToMaskBatch,
    "UR Regions To Mask Batch": URRegionsToMaskBatch,
    "UR Mask Processor": URMaskProcessor,
    "UR Keyframe Prep": URKeyframePrep,
    "UR Soft Composite": URSoftComposite,
}


NODE_DISPLAY_NAME_MAPPINGS = {
    "UR Demo Image": "UR Demo Image",
    "UR Remove Prompt Builder": "UR Remove Prompt Builder",
    "UR Region Editor": "UR Region Editor",
    "UR Boxes To Mask Batch": "UR Boxes To Mask Batch",
    "UR Regions To Mask Batch": "UR Regions To Mask Batch",
    "UR Mask Processor": "UR Mask Processor",
    "UR Keyframe Prep": "UR Keyframe Prep",
    "UR Soft Composite": "UR Soft Composite",
}