"""ComfyUI node for the P-Q-Image-2.1-Remove-Object LoRAs: draws the red outline around the mask
exactly as expected by the model, and appends the outline sentence to the instruction.

Install: copy this folder into ComfyUI/custom_nodes/ and install requirements.txt (opencv) into
ComfyUI's Python. The node appears as "Precise Removal Outline" (category Pruna).

The mask is optional. It can be painted (Load Image's MaskEditor: `mask`), come from a mask file
(e.g. Load Image (as Mask), channel red, white = object: `mask_file`), or both: they are combined
(union). Without one, or with empty ones (LoadImage gives an empty mask when nothing was painted),
the photo and the instruction pass through unchanged: the LoRA also removes objects from the
instruction alone. The outline itself comes from draw_outline.py, the same file as in
PrunaAI/P-Q-Image-2.1-Remove-Object, so it is the same ring, pixel for pixel, as in the diffusers usage.
"""

from __future__ import annotations

import numpy as np
import torch
from PIL import Image

from .draw_outline import OUTLINE_SUFFIX, RED, outline_ring


def _mask_at(mask: torch.Tensor, index: int, size: tuple[int, int]) -> np.ndarray:
    """Boolean mask of batch item `index` at the photo's size (w, h), resized and thresholded as
    draw_outline.draw_outline does (8-bit grey, nearest neighbour, > 127)."""
    m = mask if mask.dim() == 2 else mask[min(index, mask.shape[0] - 1)]
    grey = Image.fromarray(np.clip(np.rint(m.detach().cpu().float().numpy() * 255), 0, 255).astype(np.uint8), "L")
    return np.array(grey.resize(size, Image.NEAREST)) > 127


class PreciseRemovalOutline:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE", {"tooltip": "The photo to edit."}),
                "instruction": ("STRING", {"multiline": True, "default": "Remove the ",
                                           "tooltip": "What to remove, e.g. \"Remove the man on the left.\""}),
            },
            "optional": {
                "mask": ("MASK", {"tooltip": "The object to remove (paint it with the MaskEditor). Optional: "
                                             "without a mask, the instruction alone says what to remove."}),
                "mask_file": ("MASK", {"tooltip": "A mask from a file (white = the object to remove), any size: "
                                                  "resized to the photo. Combined with `mask` if both are given."}),
            },
        }

    RETURN_TYPES = ("IMAGE", "STRING")
    RETURN_NAMES = ("image", "prompt")
    OUTPUT_TOOLTIPS = ("The photo with the red outline (unchanged without a mask).",
                       "The instruction, followed by the outline sentence when an outline was drawn.")
    FUNCTION = "run"
    CATEGORY = "Pruna"
    DESCRIPTION = ("Draws the red outline expected by P-Q-Image-2.1-Remove-Object (2 px outside the "
                   "mask, 3 px wide, 4 px on photos over 900 px) and adds the matching sentence to the "
                   "instruction. Without a mask, passes both through.")

    def run(self, image: torch.Tensor, instruction: str, mask: torch.Tensor | None = None,
            mask_file: torch.Tensor | None = None):
        out, drawn = image.clone(), False
        masks = [x for x in (mask, mask_file) if x is not None and x.numel() and float(x.max()) > 0]
        if masks:
            red = torch.tensor([c / 255 for c in RED], dtype=out.dtype, device=out.device)
            for i in range(out.shape[0]):
                h, w = out.shape[1], out.shape[2]
                m = np.logical_or.reduce([_mask_at(x, i, (w, h)) for x in masks])
                if m.any():
                    ring = torch.from_numpy(outline_ring(m)).to(out.device)
                    out[i][ring] = red
                    drawn = True
        prompt = f"{instruction.rstrip()} {OUTLINE_SUFFIX}" if drawn else instruction
        return (out, prompt)


NODE_CLASS_MAPPINGS = {"PreciseRemovalOutline": PreciseRemovalOutline}
NODE_DISPLAY_NAME_MAPPINGS = {"PreciseRemovalOutline": "Precise Removal Outline"}