"""
Krea2 MultiRef custom nodes for ComfyUI.
Adapted from Ostris's implementation, with RoPE coordinate fix and
training-specific adjustments.

Training-specific details:
  - Reference tag format: <reference_N><|vision_start|><|image_pad|><|vision_end|></reference_N>
  - VLM max pixels: 512*512 = 262144 (patched from 384*384 during training)
  - VAE ref latent max pixels: 1024*1024, snapped to /16
"""

import math

import torch
from einops import rearrange

import comfy.conds
import comfy.ldm.common_dit
import comfy.utils
import node_helpers
from comfy.ldm.flux.layers import timestep_embedding

# Training-patched values
VLM_MAX_PIXELS      = 512 * 512   # patched from 384*384
REF_LATENT_MAX_PIXELS = 1024 * 1024
REF_SNAP            = 16           # VAE f8 * patch 2


# ---------------------------------------------------------------------------
# Image resizing helpers
# ---------------------------------------------------------------------------

def _fit_area(samples, max_pixels, snap=1):
    """
    Downscale (never upscale) a (B, C, H, W) tensor to fit max_pixels,
    keeping aspect ratio, snapping dims to snap.
    """
    h, w   = samples.shape[2], samples.shape[3]
    scale  = min(1.0, math.sqrt(max_pixels / (w * h)))
    nw     = max(round(w * scale / snap) * snap, snap)
    nh     = max(round(h * scale / snap) * snap, snap)
    if (nh, nw) == (h, w):
        return samples
    return comfy.utils.common_upscale(samples, nw, nh, "area", "disabled")


# ---------------------------------------------------------------------------
# Text + vision encode node
# ---------------------------------------------------------------------------

class Krea2ReferenceTextEncode:
    """
    Encodes prompt + reference images through the Krea2 Qwen3-VL text encoder.

    Reference images are injected into the prompt using the tag format
    that was used during training:
        <reference_N><|vision_start|><|image_pad|><|vision_end|></reference_N>

    The VLM sees a downscaled version (max 512*512 px) for semantic understanding.
    When a VAE is connected, full-resolution reference latents (max 1MP, snapped
    to /16) are also attached for the model patch to use.

    This node replaces the standard CLIPTextEncode when using Krea2 edit LoRAs.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "clip":   ("CLIP",),
                "prompt": ("STRING", {
                    "multiline": True,
                    "dynamicPrompts": True,
                }),
            },
            "optional": {
                "vae":    ("VAE",),
                "image1": ("IMAGE",),
                "image2": ("IMAGE",),
                "image3": ("IMAGE",),
                "image4": ("IMAGE",),
            },
        }

    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION     = "encode"
    CATEGORY     = "krea2/editing"
    DESCRIPTION  = (
        "Encode prompt + reference images for a Krea2 edit LoRA. "
        "Images are fed to the Qwen3-VL text encoder via training-matching "
        "reference tags. Connect a VAE to also attach reference latents for "
        "the Krea2ModelPatcher node."
    )

    def encode(self, clip, prompt,
               vae=None, image1=None, image2=None, image3=None, image4=None):
        images_vl   = []
        ref_latents = []
        image_prompt = ""

        for image in (image1, image2, image3, image4):
            if image is None:
                continue

            n = len(images_vl) + 1

            # Downscale for VLM (semantic understanding)
            samples_hw = image.movedim(-1, 1)    # (B, C, H, W)
            vl_img     = _fit_area(samples_hw, VLM_MAX_PIXELS).movedim(1, -1)
            images_vl.append(vl_img)

            # Build reference tag matching training format
            image_prompt += (
                f"<reference_{n}>"
                f"<|vision_start|><|image_pad|><|vision_end|>"
                f"</reference_{n}>"
            )

            # Encode reference latent for model patcher
            if vae is not None:
                s = _fit_area(samples_hw, REF_LATENT_MAX_PIXELS, snap=REF_SNAP)
                ref_latents.append(
                    vae.encode(s.movedim(1, -1)[:, :, :, :3])
                )

        # Krea2 uses its own template (KREA2_TEMPLATE) — try to import it,
        # fall back to passing None which uses the model's default.
        try:
            from comfy.text_encoders.krea2 import KREA2_TEMPLATE
            llama_template = KREA2_TEMPLATE
        except ImportError:
            llama_template = None

        tokens = clip.tokenize(
            image_prompt + prompt,
            images=images_vl if images_vl else None,
            llama_template=llama_template,
        )
        conditioning = clip.encode_from_tokens_scheduled(tokens)

        if ref_latents:
            conditioning = node_helpers.conditioning_set_values(
                conditioning,
                {"reference_latents": ref_latents},
                append=True,
            )

        return (conditioning,)


# ---------------------------------------------------------------------------
# Conditioning helper  (for the simple latent-only encoder)
# ---------------------------------------------------------------------------

def _append_ref_to_cond(conditioning, ref_tensor, strength):
    out = []
    scaled = ref_tensor * strength
    for embedding, meta in conditioning:
        meta = dict(meta)
        refs = list(meta.get("reference_latents") or [])
        refs.append(scaled)
        meta["reference_latents"] = refs
        out.append((embedding, meta))
    return out


# ---------------------------------------------------------------------------
# Reference packing  — Kontext-style RoPE layout
# ---------------------------------------------------------------------------

def _pack_refs(dit, ref_latents, bs, device, dtype, target_w):
    """
    Patchify reference latents into tokens + RoPE positions.

    RoPE layout (same axis-0=0 as target, offset in x):
      target:    axis-0=0,  y in [0, h_),    x in [0, target_w)
      ref 0:     axis-0=0,  y in [0, rh0),   x in [target_w, target_w+rw0)
      ref 1:     axis-0=0,  y in [0, rh1),   x in [target_w+rw0, ...)
    """
    patch      = dit.patch
    ref_tokens = []
    ref_pos    = []
    x_offset   = target_w

    for i, ref in enumerate(ref_latents):
        if ref.ndim == 5:
            rb, rc, rt, rh5, rw5 = ref.shape
            ref = ref.reshape(rb * rt, rc, rh5, rw5)

        ref = comfy.ldm.common_dit.pad_to_patch_size(
            ref.to(device, dtype), (patch, patch)
        )
        ref = comfy.utils.repeat_to_batch_size(ref, bs)

        rh = ref.shape[-2] // patch
        rw = ref.shape[-1] // patch

        ref_tokens.append(
            rearrange(ref, "b c (h ph) (w pw) -> b (h w) (c ph pw)",
                      ph=patch, pw=patch)
        )

        rid = torch.zeros(rh, rw, 3, device=device, dtype=torch.float32)
        rid[..., 0] = 0.0
        rid[..., 1] = torch.arange(rh, device=device,
                                   dtype=torch.float32)[:, None]
        rid[..., 2] = (
            torch.arange(rw, device=device, dtype=torch.float32)[None, :]
            + x_offset
        )
        ref_pos.append(rid.reshape(1, rh * rw, 3).repeat(bs, 1, 1))
        x_offset += rw

    return torch.cat(ref_tokens, dim=1), torch.cat(ref_pos, dim=1)


# ---------------------------------------------------------------------------
# Per-block forward with split modulation
# ---------------------------------------------------------------------------

def _block_ref_forward(block, x, vec, refvec, split, freqs, transformer_options):
    """
    SingleStreamBlock forward with per-span AdaLN:
      [:split]  text + noisy image  ->  vec    (real timestep)
      [split:]  clean references   ->  refvec  (timestep=0)
    """
    m = block.mod(vec)
    r = block.mod(refvec)

    def mod(h, scale, shift):
        return torch.cat(
            (
                (1 + m[scale]) * h[:, :split] + m[shift],
                (1 + r[scale]) * h[:, split:] + r[shift],
            ),
            dim=1,
        )

    def gate(h, g):
        return torch.cat(
            (m[g] * h[:, :split], r[g] * h[:, split:]),
            dim=1,
        )

    x = x + gate(
        block.attn(
            mod(block.prenorm(x), 0, 1),
            freqs, None,
            transformer_options=transformer_options,
        ),
        2,
    )
    x = x + gate(block.mlp(mod(block.postnorm(x), 3, 4)), 5)
    return x


# ---------------------------------------------------------------------------
# Full forward with reference latents
# ---------------------------------------------------------------------------

def _forward_with_refs(dit, x, timesteps, context, ref_latents,
                       transformer_options):
    temporal = x.ndim == 5
    if temporal:
        b5, c5, t5, h5, w5 = x.shape
        x = x.reshape(b5 * t5, c5, h5, w5)

    bs, c, H_orig, W_orig = x.shape
    patch  = dit.patch
    device = x.device

    x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch, patch))
    H, W   = x.shape[-2], x.shape[-1]
    h_, w_ = H // patch, W // patch

    context = dit._unpack_context(context)

    img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)",
                    ph=patch, pw=patch)

    reftok, refpos = _pack_refs(dit, ref_latents, bs, device, x.dtype,
                                target_w=w_)
    reflen = reftok.shape[1]

    img = dit.first(torch.cat((img, reftok), dim=1))

    t     = dit.tmlp(timestep_embedding(timesteps, dit.tdim)
                     .unsqueeze(1).to(img.dtype))
    tvec  = dit.tproj(t)
    t0    = dit.tmlp(timestep_embedding(torch.zeros_like(timesteps), dit.tdim)
                     .unsqueeze(1).to(img.dtype))
    tvec0 = dit.tproj(t0)

    context = dit.txtfusion(context, mask=None,
                            transformer_options=transformer_options)
    context = dit.txtmlp(context)

    txtlen = context.shape[1]
    imglen = img.shape[1]

    combined = torch.cat((context, img), dim=1)
    split    = txtlen + imglen - reflen

    txtpos = torch.zeros(bs, txtlen, 3, device=device, dtype=torch.float32)
    imgids = torch.zeros(h_, w_, 3, device=device, dtype=torch.float32)
    imgids[..., 1] = torch.arange(h_, device=device,
                                  dtype=torch.float32)[:, None]
    imgids[..., 2] = torch.arange(w_, device=device,
                                  dtype=torch.float32)[None, :]
    imgpos = imgids.reshape(1, h_ * w_, 3).repeat(bs, 1, 1)

    pos   = torch.cat((txtpos, imgpos, refpos), dim=1)
    freqs = dit.pe_embedder(pos)

    for block in dit.blocks:
        combined = _block_ref_forward(
            block, combined, tvec, tvec0, split, freqs, transformer_options
        )

    final = dit.last(combined, t)
    out   = final[:, txtlen:split, :]
    out   = rearrange(
        out,
        "b (h w) (c ph pw) -> b c (h ph) (w pw)",
        h=h_, w=w_, ph=patch, pw=patch, c=dit.channels,
    )
    out = out[:, :, :H_orig, :W_orig]

    if temporal:
        out = out.reshape(b5, t5, dit.channels, H_orig, W_orig).movedim(1, 2)
    return out


# ---------------------------------------------------------------------------
# Model patcher
# ---------------------------------------------------------------------------

class Krea2ModelPatcher:
    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {"model": ("MODEL",)}}

    RETURN_TYPES  = ("MODEL",)
    RETURN_NAMES  = ("model",)
    CATEGORY      = "krea2/editing"
    FUNCTION      = "patch"

    def patch(self, model):
        patched    = model.clone()
        base_model = patched.model
        dit        = patched.get_model_object("diffusion_model")

        orig_forward            = dit.forward
        orig_extra_conds        = base_model.extra_conds
        orig_extra_conds_shapes = base_model.extra_conds_shapes

        def extra_conds(**kwargs):
            out  = orig_extra_conds(**kwargs)
            refs = kwargs.get("reference_latents")
            if refs is not None:
                out["ref_latents"] = comfy.conds.CONDList(
                    [base_model.process_latent_in(lat) for lat in refs]
                )
            return out

        def extra_conds_shapes(**kwargs):
            out  = orig_extra_conds_shapes(**kwargs)
            refs = kwargs.get("reference_latents")
            if refs is not None:
                out["ref_latents"] = [
                    1, 16,
                    sum(math.prod(r.size()) for r in refs) // 16,
                ]
            return out

        def forward(
            x,
            timesteps,
            context,
            attention_mask=None,
            transformer_options={},
            ref_latents=None,
            **kwargs,
        ):
            if ref_latents is None or len(ref_latents) == 0:
                return orig_forward(
                    x, timesteps, context,
                    attention_mask=attention_mask,
                    transformer_options=transformer_options,
                    **kwargs,
                )
            return _forward_with_refs(
                dit, x, timesteps, context, ref_latents, transformer_options,
            )

        patched.add_object_patch("extra_conds",             extra_conds)
        patched.add_object_patch("extra_conds_shapes",      extra_conds_shapes)
        patched.add_object_patch("diffusion_model.forward", forward)
        return (patched,)


# ---------------------------------------------------------------------------
# Simple latent-only reference encoder  (no VLM — use when you already have
# a text conditioning from a normal CLIPTextEncode and just want to attach
# latent refs without re-encoding text)
# ---------------------------------------------------------------------------

class Krea2ReferenceEncoder:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "positive":    ("CONDITIONING",),
                "base_latent": ("LATENT", {
                    "tooltip": "VAE-encoded image to edit. Reference slot 1."
                }),
                "strength": ("FLOAT", {
                    "default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05,
                }),
            },
            "optional": {
                "reference_latent_2": ("LATENT",),
                "reference_latent_3": ("LATENT",),
                "reference_latent_4": ("LATENT",),
            },
        }

    RETURN_TYPES = ("CONDITIONING",)
    CATEGORY     = "krea2/editing"
    FUNCTION     = "encode"

    def encode(
        self,
        positive,
        base_latent,
        strength           = 1.0,
        reference_latent_2 = None,
        reference_latent_3 = None,
        reference_latent_4 = None,
    ):
        slots   = [base_latent, reference_latent_2,
                   reference_latent_3, reference_latent_4]
        tensors = [s["samples"] for s in slots if s is not None]

        new_positive = positive
        for ref in tensors:
            new_positive = _append_ref_to_cond(new_positive, ref, strength)
        return (new_positive,)


# ---------------------------------------------------------------------------
# Diagnostic
# ---------------------------------------------------------------------------

class Krea2DiagnosticNode:
    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {"model": ("MODEL",)}}

    RETURN_TYPES = ()
    OUTPUT_NODE  = True
    CATEGORY     = "krea2/editing"
    FUNCTION     = "run"

    def run(self, model):
        import inspect
        dit = model.get_model_object("diffusion_model")
        print("\n=== DIT TYPE ===")
        print(type(dit).__name__, type(dit).__module__)
        print("\n=== DIT NAMED CHILDREN ===")
        for name, mod in dit.named_children():
            print(f"  {name}: {type(mod).__name__}")
        print("\n=== DIT NON-MODULE ATTRS ===")
        for k, v in vars(dit).items():
            if not k.startswith("_") and not isinstance(v, torch.nn.Module):
                print(f"  {k}: {type(v).__name__} = {repr(v)[:120]}")
        print("\n=== DIT FORWARD SIGNATURE ===")
        print(inspect.signature(dit.forward))
        print("\n=== FIRST BLOCK ===")
        if hasattr(dit, "blocks"):
            block = dit.blocks[0]
            print(f"  {type(block).__name__}")
            for n, m in block.named_children():
                print(f"    {n}: {type(m).__name__}")
            print(f"  sig: {inspect.signature(block.forward)}")
        print(f"\n  has _unpack_context: {hasattr(dit, '_unpack_context')}")
        return ()


# ---------------------------------------------------------------------------
# Registration
# ---------------------------------------------------------------------------

NODE_CLASS_MAPPINGS = {
    "Krea2ModelPatcher":          Krea2ModelPatcher,
    "Krea2ReferenceEncoder":      Krea2ReferenceEncoder,
    "Krea2ReferenceTextEncode":   Krea2ReferenceTextEncode,
    "Krea2DiagnosticNode":        Krea2DiagnosticNode,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "Krea2ModelPatcher":          "Krea2 Model Patcher",
    "Krea2ReferenceEncoder":      "Krea2 Reference Encoder (latent only)",
    "Krea2ReferenceTextEncode":   "Krea2 Reference Text Encode",
    "Krea2DiagnosticNode":        "Krea2 Diagnostic",
}