import torch

class RGBA_Black_Background:
    """
    Composites an RGBA image over a black background and outputs RGB.
    Assumes input image in ComfyUI tensor format: float32, 0..1, shape [B,H,W,C].
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),  # Can be RGBA or RGB; if RGB, it passes through
            }
        }

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

    def composite(self, image):
        # image: [B,H,W,C], float32 0..1
        if not isinstance(image, torch.Tensor):
            raise TypeError("Expected ComfyUI IMAGE as a torch.Tensor")

        if image.ndim != 4:
            raise ValueError(f"Expected image shape [B,H,W,C], got {tuple(image.shape)}")

        b, h, w, c = image.shape

        # If already RGB, just ensure it's sane and return.
        if c == 3:
            return (image.clamp(0.0, 1.0),)

        if c != 4:
            raise ValueError(f"Expected 3 (RGB) or 4 (RGBA) channels, got {c}")

        rgb = image[..., :3]
        a   = image[..., 3:4]  # [B,H,W,1]

        # Black background composite: bg=0 => out = rgb*a + 0*(1-a) = rgb*a
        out = rgb * a

        return (out.clamp(0.0, 1.0),)


NODE_CLASS_MAPPINGS = {
    "RGBA_Black_Background": RGBA_Black_Background
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "RGBA_Black_Background": "RGBA to Black Background RGB (Black Background)"
}