# paste_rgba_at_xy.py
# ComfyUI custom node: Paste a small RGBA image onto a larger RGBA canvas at (x, y)
# Supports alpha compositing ("source-over") or hard replace.
#
# Usage:
#  - canvas: big RGBA image (IMAGE)
#  - overlay: small RGBA image (IMAGE)
#  - x, y: top-left destination coordinate on the canvas
#  - blend_mode: "alpha_over" (default) or "replace"
#  - If either input is RGB (3ch), it will be upgraded to RGBA with alpha=1.0
#  - Batching:
#       * If one input has batch size 1 and the other >1, the single image is broadcast.
#       * If both have batch >1, their batch sizes must match.

import torch

class PasteRGBAAtXY:
    """
    Paste a small RGBA image on a larger RGBA canvas at a specified (x, y) coordinate.

    - Default blend mode is proper alpha compositing (SRC over).
    - "replace" mode copies the overlay's RGBA pixels directly (no blending).
    - Handles out-of-bounds and negative coordinates by clipping.
    - Works with batches. Broadcasts a single overlay across a batch of canvases (and vice versa).
    """

    CATEGORY = "image/compose"
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "paste"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "canvas": ("IMAGE",),
                "overlay": ("IMAGE",),
                "x": ("INT", {"default": 0, "min": -32768, "max": 32768, "step": 1}),
                "y": ("INT", {"default": 0, "min": -32768, "max": 32768, "step": 1}),
                "blend_mode": (["alpha_over", "replace"], {"default": "alpha_over"}),
            }
        }

    @staticmethod
    def _ensure_rgba(img: torch.Tensor) -> torch.Tensor:
        """
        Ensure an IMAGE tensor is RGBA. If RGB, append opaque alpha.
        Shape convention in ComfyUI for IMAGE is [B, H, W, C] with float32 in [0, 1].
        """
        if img.ndim != 4:
            raise ValueError(f"Expected image tensor with 4 dims [B,H,W,C], got shape {tuple(img.shape)}")
        c = img.shape[-1]
        if c == 4:
            return img
        if c == 3:
            alpha = torch.ones((*img.shape[:-1], 1), dtype=img.dtype, device=img.device)
            return torch.cat([img, alpha], dim=-1)
        raise ValueError(f"Expected 3 or 4 channels, got {c}")

    @staticmethod
    def _pair_count(b1: int, b2: int) -> int:
        """Return the output batch size if broadcasting is allowed."""
        if b1 == b2:
            return b1
        if b1 == 1 or b2 == 1:
            return max(b1, b2)
        raise ValueError(f"Incompatible batch sizes: {b1} vs {b2}")

    @staticmethod
    def _get_batch(img: torch.Tensor, i: int) -> torch.Tensor:
        """Fetch the i-th batch image with broadcasting if needed."""
        b = img.shape[0]
        return img[0] if b == 1 else img[i]

    @staticmethod
    def _alpha_over(dst_rgba: torch.Tensor, src_rgba: torch.Tensor, dx: int, dy: int) -> None:
        """
        Alpha-composite src onto dst in-place at integer offset (dx, dy).
        Both tensors are HxWx4, float in [0,1]. Clips automatically when out of bounds.
        """
        Hc, Wc, _ = dst_rgba.shape
        Ho, Wo, _ = src_rgba.shape

        # Compute intersection rectangle on canvas (destination)
        x0 = max(0, dx)
        y0 = max(0, dy)
        x1 = min(Wc, dx + Wo)
        y1 = min(Hc, dy + Ho)

        if x1 <= x0 or y1 <= y0:
            return  # Nothing overlaps

        # Corresponding source crop
        sx0 = x0 - dx
        sy0 = y0 - dy
        w = x1 - x0
        h = y1 - y0

        dst_region = dst_rgba[y0:y0+h, x0:x0+w, :]
        src_region = src_rgba[sy0:sy0+h, sx0:sx0+w, :]

        # Split channels
        cb = dst_region[..., :3]
        ab = dst_region[..., 3:4]
        co = src_region[..., :3]
        ao = src_region[..., 3:4]

        # Premultiply colors
        cb_p = cb * ab
        co_p = co * ao

        # Source-over composition
        out_a = ao + ab * (1.0 - ao)
        out_c_p = co_p + cb_p * (1.0 - ao)

        # Convert back to straight (guard against divide by zero)
        eps = 1e-8
        out_c = torch.where(out_a > eps, out_c_p / out_a.clamp_min(eps), torch.zeros_like(out_c_p))

        # Write back (clamp just in case)
        dst_region[..., :3] = out_c.clamp(0.0, 1.0)
        dst_region[..., 3:4] = out_a.clamp(0.0, 1.0)

        dst_rgba[y0:y0+h, x0:x0+w, :] = dst_region

    @staticmethod
    def _replace(dst_rgba: torch.Tensor, src_rgba: torch.Tensor, dx: int, dy: int) -> None:
        """
        Direct overwrite (copy) of src_rgba into dst_rgba at (dx, dy), clipped to bounds.
        """
        Hc, Wc, _ = dst_rgba.shape
        Ho, Wo, _ = src_rgba.shape

        x0 = max(0, dx)
        y0 = max(0, dy)
        x1 = min(Wc, dx + Wo)
        y1 = min(Hc, dy + Ho)
        if x1 <= x0 or y1 <= y0:
            return

        sx0 = x0 - dx
        sy0 = y0 - dy
        w = x1 - x0
        h = y1 - y0

        dst_rgba[y0:y0+h, x0:x0+w, :] = src_rgba[sy0:sy0+h, sx0:sx0+w, :]

    def paste(self, canvas, overlay, x, y, blend_mode):
        # Ensure RGBA
        canvas = self._ensure_rgba(canvas)
        overlay = self._ensure_rgba(overlay)

        b_canvas = canvas.shape[0]
        b_overlay = overlay.shape[0]
        out_b = self._pair_count(b_canvas, b_overlay)

        out_list = []
        for i in range(out_b):
            base = self._get_batch(canvas, i).clone()  # HxWx4
            over = self._get_batch(overlay, i)         # HxWx4

            if blend_mode == "alpha_over":
                self._alpha_over(base, over, int(x), int(y))
            else:  # "replace"
                self._replace(base, over, int(x), int(y))

            out_list.append(base)

        out = torch.stack(out_list, dim=0)
        return (out,)


# --- ComfyUI registration ---
NODE_CLASS_MAPPINGS = {
    "PasteRGBAAtXY": PasteRGBAAtXY,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "PasteRGBAAtXY": "Paste RGBA at (X,Y)",
}