"""Krea 2 reference-image conditioning for ComfyUI.

Krea 2 conditions on 12 tapped hidden states of Qwen3-VL-4B. Qwen3-VL is a
vision-language model, so reference images can be interleaved into the prompt
as native vision tokens and the DiT reads character/style/pose information
from the same hidden states it was trained on. The Krea 2 DiT has no
reference-latent pathway (unlike Flux Kontext / Qwen-Image-Edit), so no VAE
input is needed here.
"""

import re

import torch

import comfy.utils
import node_helpers

# Longest-side pixel budget fed to the vision tower. More pixels = more vision
# tokens = stronger, more literal influence from the reference.
DETAIL = {"low": 256, "normal": 512, "high": 1024, "max": 1280}
DETAIL_TIP = ("Longest side this reference is scaled to before the vision encoder. Higher = more "
              "vision tokens = stronger, more literal reference. Use high for the character/subject "
              "and low/normal for pose or style hints to avoid side-by-side collages.")
SPLIT_TIP = ("Fraction of the denoise steps that see TEXT-ONLY conditioning before the references "
             "kick in. The first steps fix the layout, so 0.25 strongly prevents side-by-side "
             "collages of the references. 0 disables.")

# Vision tower consumes 32px merged patches; align shuffle tiles to that grid.
VISION_PATCH = 32


MASK_TIP = ("Optional subject mask (e.g. from RemoveBackground). The reference is cropped to the "
            "mask's bounding box and the background is whited out before the vision encoder, so "
            "messy references (scenes, watermarks, other characters) only contribute the subject.")
BG_TIP = ("Optional: connect Load Background Removal Model (birefnet) and every content-role "
          "reference is auto-cleaned — cropped to its subject with the background whited out. "
          "Per-image masks override it; style-role images are left untouched (their background "
          "often carries the style).")


def isolate_subject(s, mask, pad=0.05):
    """Crop NCHW image to the mask bbox (+pad) and white out the background."""
    if mask.shape[-2:] != s.shape[-2:]:
        mask = comfy.utils.common_upscale(mask.reshape(-1, 1, *mask.shape[-2:]), s.shape[3], s.shape[2], "bilinear", "disabled").reshape(-1, *s.shape[-2:])
    m = mask[0] > 0.5
    if not m.any():
        return s
    ys, xs = torch.where(m)
    y0, y1 = ys.min().item(), ys.max().item() + 1
    x0, x1 = xs.min().item(), xs.max().item() + 1
    ph, pw = round((y1 - y0) * pad), round((x1 - x0) * pad)
    y0, y1 = max(0, y0 - ph), min(s.shape[2], y1 + ph)
    x0, x1 = max(0, x0 - pw), min(s.shape[3], x1 + pw)
    mm = mask[:1].unsqueeze(1).to(s.dtype)
    s = s * mm + (1.0 - mm)
    return s[:, :, y0:y1, x0:x1]


def shuffle_patches(s, tile):
    """Randomly permute tile positions (NCHW). Deterministic. Destroys global
    shape while keeping patch-local style statistics."""
    n, c, h, w = s.shape
    gh, gw = max(1, h // tile), max(1, w // tile)
    s = s[:, :, :gh * tile, :gw * tile]
    tiles = s.reshape(n, c, gh, tile, gw, tile).permute(0, 2, 4, 1, 3, 5).reshape(n, gh * gw, c, tile, tile)
    perm = torch.randperm(gh * gw, generator=torch.Generator().manual_seed(0))
    tiles = tiles[:, perm]
    return tiles.reshape(n, gh, gw, c, tile, tile).permute(0, 3, 1, 4, 2, 5).reshape(n, c, gh * tile, gw * tile)

# Same edit template Qwen-Image-Edit uses; Krea 2 was trained on Qwen3-VL
# hidden states over this chat format, and the Krea2 text encoder strips the
# system + user-opening prefix before conditioning.
EDIT_TEMPLATE = (
    "<|im_start|>system\n"
    "Describe the key features of the input image (color, shape, size, texture, "
    "objects, background), then explain how the user's text instruction should "
    "alter or modify the image. Generate a new image that meets the user's "
    "requirements while maintaining consistency with the original input where "
    "appropriate.<|im_end|>\n"
    "<|im_start|>user\n{}<|im_end|>\n"
    "<|im_start|>assistant\n"
)

# "image1", "image 2", "img3", "imagen 1" -> "Picture N" (what the encoder sees)
_IMG_WORD = re.compile(r"\b(?:image|img|imagen|picture)[ _]?([1-4])\b", re.IGNORECASE)

# Role-typed shorthand: "character1" / "style2" expand to the exact clauses the
# reference LoRA trains on. Roles are prompt-driven: styleN shuffles that slot's
# patches (style-only transfer); characterN keeps it literal and wins over styleN.
_CHAR_WORD = re.compile(r"\b(?:character|char|personaje)[ _]?([1-4])\b", re.IGNORECASE)
_STYLE_WORD = re.compile(r"\b(?:style|estilo)[ _]?([1-4])\b", re.IGNORECASE)


def expand_role_tokens(prompt):
    """'The character1 in a car in the style2' ->
    'The character from Picture 1 in a car in the exact art style of Picture 2'.
    Returns (text, roles): roles are prompt-driven — styleN marks the slot as
    style (patch shuffle); characterN marks it content and wins if both appear
    for the same slot (identity needs the unshuffled image)."""
    style_slots = {int(m.group(1)) - 1 for m in _STYLE_WORD.finditer(prompt)}
    char_slots = {int(m.group(1)) - 1 for m in _CHAR_WORD.finditer(prompt)}
    text = _CHAR_WORD.sub(lambda m: "character from Picture " + m.group(1), prompt)
    text = _STYLE_WORD.sub(lambda m: "exact art style of Picture " + m.group(1), text)
    roles = tuple("style" if i in style_slots and i not in char_slots else "content"
                  for i in range(4))
    return text, roles


# Auto-added when the prompt doesn't mention a connected image. Mirrors the
# phrasing the reference LoRA trains on, so casual prompts still bind the refs.
ROLE_CLAUSE = {
    "content": "the character from Picture {}",
    "style": "in the exact art style of Picture {}",
}


def build_image_prompt(prompt, image_slots, roles=None):
    """Returns (full_prompt, connected_slot_indices). Slot numbering is
    preserved so 'image2' still means the image2 input even if image1 is empty.
    Connected images never mentioned in the prompt get a role clause appended."""
    image_prompt = ""
    connected = []
    for i, image in enumerate(image_slots):
        if image is None:
            continue
        connected.append(i)
        image_prompt += "Picture {}: <|vision_start|><|image_pad|><|vision_end|>".format(i + 1)
    text = _IMG_WORD.sub(lambda m: "Picture " + m.group(1), prompt)
    if roles is not None:
        missing = [i for i in connected if "Picture {}".format(i + 1) not in text]
        clauses = [ROLE_CLAUSE[roles[i]].format(i + 1) for i in missing]
        if clauses:
            text = (text.rstrip().rstrip(".,") + ", " if text.strip() else "") + ", ".join(clauses)
    return image_prompt + text, connected


class TextEncodeKrea2Reference:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "clip": ("CLIP",),
                "prompt": ("STRING", {
                    "multiline": True,
                    "dynamicPrompts": True,
                    "tooltip": "Roles come from the prompt: character1..4 = identity reference, style1..4 = style-only (patches shuffled). Both expand to the clauses the LoRA trains on. Plain image1..4 / Picture 1..4 also work, "
                               "e.g. 'Redraw the character from image1 in the pose of image2'.",
                }),
            },
            "optional": {
                "image1": ("IMAGE",),
                "image1_detail": (list(DETAIL), {"default": "high", "tooltip": DETAIL_TIP}),
                "image1_mask": ("MASK", {"tooltip": MASK_TIP}),
                "image2": ("IMAGE",),
                "image2_detail": (list(DETAIL), {"default": "high", "tooltip": DETAIL_TIP}),
                "image2_mask": ("MASK", {"tooltip": MASK_TIP}),
                "image3": ("IMAGE",),
                "image3_detail": (list(DETAIL), {"default": "high", "tooltip": DETAIL_TIP}),
                "image3_mask": ("MASK", {"tooltip": MASK_TIP}),
                "image4": ("IMAGE",),
                "image4_detail": (list(DETAIL), {"default": "high", "tooltip": DETAIL_TIP}),
                "image4_mask": ("MASK", {"tooltip": MASK_TIP}),
                "bg_removal": ("BACKGROUND_REMOVAL", {"tooltip": BG_TIP}),
                "layout_split": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 0.9, "step": 0.05, "tooltip": SPLIT_TIP}),
            },
        }

    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "encode"
    CATEGORY = "conditioning/krea2"
    DESCRIPTION = ("Encodes a prompt plus up to 4 reference images for Krea 2. References enter "
                   "as Qwen3-VL vision tokens, enabling character/style/pose transfer "
                   "('Redraw the character from image1 in the pose of image2').")

    def encode(self, clip, prompt, image1=None, image2=None, image3=None, image4=None,
               image1_detail="high", image2_detail="high", image3_detail="high", image4_detail="high",
               image1_mask=None, image2_mask=None, image3_mask=None, image4_mask=None,
               bg_removal=None, layout_split=0.0):
        slots = (image1, image2, image3, image4)
        details = (image1_detail, image2_detail, image3_detail, image4_detail)
        masks = (image1_mask, image2_mask, image3_mask, image4_mask)
        prompt, roles = expand_role_tokens(prompt)
        full_prompt, connected = build_image_prompt(prompt, slots, roles)

        images_vl = []
        for i in connected:
            target = DETAIL[details[i]]
            s = slots[i].movedim(-1, 1)  # NHWC -> NCHW
            mask = masks[i]
            if mask is None and bg_removal is not None and roles[i] == "content":
                mask = bg_removal.encode_image(slots[i][:, :, :, :3])
            if mask is not None:
                s = isolate_subject(s, mask)
            scale = target / max(s.shape[2], s.shape[3])
            w = max(VISION_PATCH, round(s.shape[3] * scale))
            h = max(VISION_PATCH, round(s.shape[2] * scale))
            s = comfy.utils.common_upscale(s, w, h, "area", "disabled")
            if roles[i] == "style":
                s = shuffle_patches(s, 2 * VISION_PATCH if target >= 512 else VISION_PATCH)
            images_vl.append(s.movedim(1, -1)[:, :, :, :3])

        tokens = clip.tokenize(
            full_prompt,
            images=images_vl,
            llama_template=EDIT_TEMPLATE if images_vl else None,
        )
        cond = clip.encode_from_tokens_scheduled(tokens)

        if images_vl and layout_split > 0.0:
            # First steps see text-only conditioning: layout gets fixed before the
            # reference tokens can push it toward a side-by-side collage.
            text_tokens = clip.tokenize(_IMG_WORD.sub(lambda m: "Picture " + m.group(1), prompt))
            text_cond = clip.encode_from_tokens_scheduled(text_tokens)
            cond = (node_helpers.conditioning_set_values(text_cond, {"start_percent": 0.0, "end_percent": layout_split})
                    + node_helpers.conditioning_set_values(cond, {"start_percent": layout_split, "end_percent": 1.0}))
        return (cond,)


class Krea2ReferenceLoraLoader:
    """Load the reference LoRA with independent strengths per layer group.

    The reference addressing and anti-bleed behavior live in text_fusion/txt_in
    (the conditioning path); the DiT body tensors mostly carry identity drift.
    strength_text 1.0 + strength_dit 0.0 keeps the base model's rendering while
    following references; raise strength_dit for extra scene adherence."""

    @classmethod
    def INPUT_TYPES(cls):
        import folder_paths
        return {
            "required": {
                "model": ("MODEL",),
                "lora_name": (folder_paths.get_filename_list("loras"),),
                "strength_text_fusion": ("FLOAT", {"default": 1.0, "min": -2.0, "max": 2.0, "step": 0.05,
                                                   "tooltip": "text_fusion + txt_in: reference addressing / anti-bleed / style transfer."}),
                "strength_dit_body": ("FLOAT", {"default": 0.0, "min": -2.0, "max": 2.0, "step": 0.05,
                                                "tooltip": "DiT blocks: extra scene adherence, but drifts identity above ~0.3."}),
            }
        }

    RETURN_TYPES = ("MODEL",)
    FUNCTION = "load"
    CATEGORY = "loaders"

    def load(self, model, lora_name, strength_text_fusion, strength_dit_body):
        import comfy.sd
        import folder_paths
        lora = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("loras", lora_name), safe_load=True)
        text_keys = ("transformer.text_fusion.", "transformer.txt_in.")
        groups = (
            ({k: v for k, v in lora.items() if k.startswith(text_keys)}, strength_text_fusion),
            ({k: v for k, v in lora.items() if not k.startswith(text_keys)}, strength_dit_body),
        )
        for sd, strength in groups:
            if sd and strength != 0.0:
                model, _ = comfy.sd.load_lora_for_models(model, None, sd, strength, 0.0)
        return (model,)


NODE_CLASS_MAPPINGS = {
    "TextEncodeKrea2Reference": TextEncodeKrea2Reference,
    "Krea2ReferenceLoraLoader": Krea2ReferenceLoraLoader,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "TextEncodeKrea2Reference": "Krea 2 Reference Encode",
    "Krea2ReferenceLoraLoader": "Krea 2 Reference LoRA (layered)",
}