"""ComfyUI nodes for the native MS-Diffusion chain.

Intended graph (everything except the three MS nodes is a stock ComfyUI node):

  CheckpointLoaderSimple -> MODEL / CLIP / VAE
  LoadImage xN -> native ImageBatch (chained N-1 times) -> CLIPVisionEncode
                                                      -> CLIP_VISION_OUTPUT
  CLIPTextEncode x2 -> positive / negative CONDITIONING
  MSDiffusionLoader ------------------------------> MS_ADAPTER
  MSDiffusionConditioning ------------------------> MS_EMBEDS (+ clean prompt)
  MSDiffusionApply(MODEL, MS_ADAPTER, MS_EMBEDS) --> patched MODEL
  EmptyLatentImage -> KSampler -> VAEDecode -> SaveImage

No diffusers / transformers / omegaconf / einops / PIL imports anywhere.
"""

import os
import re

import torch
import folder_paths

from .core import (MSAdapter, MSEmbeds, adapter_hidden_sizes, build_block_weights,
                   enumerate_attn2, make_attn2_patch, new_state, NUM_TOKENS,
                   NUM_DUMMY_TOKENS)

ADAPTER_FOLDER = "ms_adapter"
_ADAPTER_EXTS = (".bin", ".safetensors", ".pt", ".pth", ".ckpt")


def _register_folders():
    """Register models/ipadapter and models/photomaker as one virtual folder.

    The upstream plugin hard-coded models/photomaker; the user's actual file is
    in models/ipadapter. Support both without touching ComfyUI source.
    """
    try:
        if ADAPTER_FOLDER not in folder_paths.folder_names_and_paths:
            folder_paths.folder_names_and_paths[ADAPTER_FOLDER] = ([], set())
        for sub in ("ipadapter", "photomaker"):
            p = os.path.join(folder_paths.models_dir, sub)
            if os.path.isdir(p):
                folder_paths.add_model_folder_path(ADAPTER_FOLDER, p)
    except Exception:
        pass


_register_folders()


def _adapter_files():
    try:
        files = [f for f in folder_paths.get_filename_list(ADAPTER_FOLDER)
                 if f.lower().endswith(_ADAPTER_EXTS)]
    except Exception:
        files = []
    if not files:
        # fallback scan, keeps the node usable even if registration was blocked
        for sub in ("ipadapter", "photomaker"):
            d = os.path.join(folder_paths.models_dir, sub)
            if os.path.isdir(d):
                files += [sub + "/" + f for f in sorted(os.listdir(d))
                          if f.lower().endswith(_ADAPTER_EXTS)]
    files = sorted(set(files)) or ["none"]
    # models/ipadapter often holds many unrelated IP-Adapter weights; put the
    # real MS-Diffusion checkpoint first so the default is the correct file.
    preferred = [f for f in files if os.path.basename(f).lower().startswith("ms_adapter")]
    return preferred + [f for f in files if f not in preferred]


def _adapter_path(name):
    try:
        return folder_paths.get_full_path_or_raise(ADAPTER_FOLDER, name)
    except Exception:
        return os.path.join(folder_paths.models_dir, name.replace("/", os.sep))


# ---------------------------------------------------------------------------
# phrase token position (port of msdiffusion/utils.py, no transformers import)
# ---------------------------------------------------------------------------

def _token_encoder(clip):
    """Return f(word) -> len(HF tokenizer.encode(word)).

    ComfyUI already owns a CLIP tokenizer instance; we just borrow it, so this
    plugin adds no transformers import of its own.
    """
    try:
        tk = clip.tokenizer.clip_l.tokenizer
        return lambda w: len(tk.encode(w))
    except Exception:
        pass
    # fallback: derive from ComfyUI's own token/weight pairs
    def from_comfy(w):
        pairs = clip.tokenize(w)["l"][0]
        n = 0
        for tid, _ in pairs:
            if int(tid) == 49407:      # EOS == pad token
                break
            n += 1
        return max(n, 1) + 1          # [BOS] + word tokens => n-1 + 2
    return from_comfy


def _is_equal_words(pr_words, ph_words):
    if len(pr_words) != len(ph_words):
        return False
    for pr, ph in zip(pr_words, ph_words):
        if "-" + ph not in pr and ph != re.sub(r'[.!?,:]$', '', pr):
            return False
    return True


def _get_phrase_idx(enc_len, phrase, prompt, num=0):
    phrase_words = phrase.split()
    if len(phrase_words) == 0:
        return [0, 0]
    prompt_words = prompt.split()
    start, end = 1, 0
    for i in range(len(prompt_words)):
        if _is_equal_words(prompt_words[i:i + len(phrase_words)], phrase_words):
            if num != 0:
                num -= 1
                continue
            end = start
            prompt_words[i + len(phrase_words) - 1] = re.sub(
                r'[.!?,:]$', '', prompt_words[i + len(phrase_words) - 1])
            for j in range(i, i + len(phrase_words)):
                end += enc_len(prompt_words[j]) - 2
            break
        else:
            start += enc_len(prompt_words[i]) - 2
    if end == 0:
        return [0, 0]
    return [start, end]


def _parse_boxes(text, num_ref, enabled):
    """Upstream: mask_threshold == 0 disables layout guidance -> all-zero boxes."""
    zeros = [[0.0, 0.0, 0.0, 0.0] for _ in range(num_ref)]
    if not enabled or not text.strip():
        return zeros
    groups = [g.strip() for g in text.strip().split(";") if g.strip()]
    parsed = []
    for g in groups:
        vals = [float(x.strip()) for x in g.split(",")]
        if len(vals) != 4:
            raise ValueError("MS-Diffusion: each box must have 4 floats (x1,y1,x2,y2); got %r" % g)
        parsed.append(vals)
    if len(parsed) < num_ref:
        # A zero box masks that subject's 16 image tokens completely (-10000 bias),
        # silently dropping it. Upstream raises here instead of guessing; do the same.
        raise ValueError(
            "MS-Diffusion: %d subjects but only %d boxes. Provide one "
            "'x1,y1,x2,y2' per subject separated by ';' (or set mask_threshold=0 "
            "to disable layout guidance)." % (num_ref, len(parsed)))
    return parsed[:num_ref]


def _clip_l_pooled(clip, phrase):
    """CLIP-L 768-d pooled output (the grounding phrase embedding).

    SDXL's CONDITIONING pooled_output is clip_g (1280) and cannot be used here.
    """
    clip.load_model()
    dev = clip.patcher.load_device
    try:
        clip.cond_stage_model.clip_l.set_clip_options({"execution_device": dev})
    except Exception:
        pass
    tokens = clip.tokenize(phrase)["l"]
    with torch.no_grad():
        out = clip.cond_stage_model.clip_l.encode_token_weights(tokens)
    pooled = out[1]
    return pooled.to(torch.float32).cpu()


# ---------------------------------------------------------------------------
# Nodes
# ---------------------------------------------------------------------------

class MSDiffusionLoader:
    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {"ms_adapter": (_adapter_files(),)}}

    RETURN_TYPES = ("MS_ADAPTER",)
    RETURN_NAMES = ("ms_adapter",)
    FUNCTION = "load"
    CATEGORY = "MSdiffusion/Native"
    DESCRIPTION = ("Load ms_adapter.bin from models/ipadapter or models/photomaker. "
                   "Pure torch; no diffusers/transformers needed.")

    def load(self, ms_adapter):
        if ms_adapter == "none":
            raise RuntimeError("MS-Diffusion: no ms_adapter.bin found in models/ipadapter or models/photomaker")
        path = _adapter_path(ms_adapter)
        sizes = adapter_hidden_sizes(path)
        adapter = MSAdapter(sizes).load(path)
        print("[MS-Diffusion] loaded %s (%d cross-attn groups)" % (path, len(sizes)))
        return (adapter,)


class MSDiffusionConditioning:
    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {
            "ms_adapter": ("MS_ADAPTER",),
            "clip": ("CLIP",),
            "clip_vision_output": ("CLIP_VISION_OUTPUT",),
            "prompt": ("STRING", {"multiline": True,
                                  "default": "best quality, high quality, a [dog] on the beach"}),
            "boxes": ("STRING", {"multiline": True,
                                 "default": "0.25, 0.25, 0.75, 0.75",
                                 "tooltip": "one 'x1,y1,x2,y2' per subject, separated by ';'. "
                                            "Ignored when mask_threshold is 0."}),
            "mask_threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05,
                                         "tooltip": "0 disables layout guidance entirely (upstream behaviour). "
                                                    ">0 enables box guidance + pseudo attention mask."}),
            "drop_grounding_tokens": ("BOOLEAN", {"default": False}),
        }}

    RETURN_TYPES = ("MS_EMBEDS", "STRING")
    RETURN_NAMES = ("ms_embeds", "prompt_clean")
    FUNCTION = "encode"
    CATEGORY = "MSdiffusion/Native"
    DESCRIPTION = ("Build MS-Diffusion image prompt tokens from CLIP-Vision penultimate "
                   "hidden states + CLIP-L phrase pooled embeddings + layout boxes. "
                   "Connect prompt_clean to your positive CLIPTextEncode so the two match.")

    def encode(self, ms_adapter, clip, clip_vision_output, prompt, boxes,
               mask_threshold, drop_grounding_tokens):
        import comfy.model_management as mm

        phrases = [p.strip() for p in re.findall(r'\[(.*?)\]', prompt)]
        if not phrases:
            raise ValueError("MS-Diffusion: subject names must be enclosed in [ ] inside the prompt")
        for p in phrases:
            if " " in p:
                raise ValueError("MS-Diffusion: subject %r must be a single word" % p)
        phrases = list(dict.fromkeys(phrases))
        clean_prompt = prompt.replace("]", " ").replace("[", " ")

        hidden = clip_vision_output["penultimate_hidden_states"]
        if hidden.dim() != 3:
            raise ValueError("MS-Diffusion: CLIP-Vision output must be penultimate_hidden_states [B,257,1664]")
        num_ref = hidden.shape[0]
        if num_ref != len(phrases):
            raise ValueError(
                "MS-Diffusion: got %d subject images but %d [subject] names in prompt (%s)"
                % (num_ref, len(phrases), phrases))
        if hidden.shape[-1] != 1664:
            raise ValueError("MS-Diffusion: expected CLIP-ViT-bigG hidden size 1664, got %d. "
                             "Use clip_vision_g.safetensors." % hidden.shape[-1])

        layout_on = mask_threshold > 0.0
        box_vals = _parse_boxes(boxes, num_ref, layout_on)
        boxes_t = torch.tensor([box_vals], dtype=torch.float32)          # [1, num_ref, 4]
        batch_boxes = boxes_t.view(num_ref, 4)                           # Resampler grounding input

        phrase_embeds = torch.cat([_clip_l_pooled(clip, p) for p in phrases], dim=0)
        if phrase_embeds.shape != (num_ref, 768):
            raise RuntimeError("MS-Diffusion: CLIP-L pooled shape %s, expected (%d,768)"
                               % (tuple(phrase_embeds.shape), num_ref))

        enc_len = _token_encoder(clip)
        seen = {}
        pidx = []
        for p in phrases:
            n = seen.get(p, 0)
            seen[p] = n + 1
            pidx.append(_get_phrase_idx(enc_len, p, clean_prompt, num=n))
        phrase_idxes = torch.tensor([pidx], dtype=torch.int32)           # [1, num_ref, 2]

        dev = mm.get_torch_device()
        dtype = torch.float16
        with torch.no_grad():
            ms_adapter.image_proj.to(device=dev, dtype=dtype)
            x = hidden.to(device=dev, dtype=dtype)
            grounding = {
                "boxes": batch_boxes.to(device=dev, dtype=dtype),
                "phrase_embeds": phrase_embeds.to(device=dev, dtype=dtype),
                "drop_grounding_tokens": [1 if drop_grounding_tokens else 0],
            }
            ip_tokens = ms_adapter.image_proj(x, grounding_kwargs=grounding)
            ip_tokens = ip_tokens.view(1, -1, ip_tokens.shape[-2], ip_tokens.shape[-1])
            ip_tokens = ip_tokens.view(1, ip_tokens.shape[-3] * ip_tokens.shape[-2], ip_tokens.shape[-1])
            dummy = ms_adapter.dummy_image_tokens.to(device=dev, dtype=dtype)
            image_prompt_embeds = torch.cat([dummy, ip_tokens], dim=1).float().cpu()
            ms_adapter.image_proj.to(device=torch.device("cpu"))
            ms_adapter._placed = None

        assert image_prompt_embeds.shape == (1, NUM_DUMMY_TOKENS + NUM_TOKENS * num_ref, 2048)
        embeds = MSEmbeds(
            image_prompt_embeds=image_prompt_embeds,
            boxes=boxes_t,
            phrase_idxes=phrase_idxes,
            num_ref=num_ref,
            scale=1.0,                                   # set by MSDiffusionApply
            mask_threshold=None if not layout_on else float(mask_threshold),
            start_step=5,                               # set by MSDiffusionApply
        )
        return (embeds, clean_prompt)


class MSDiffusionApply:
    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {
            "model": ("MODEL",),
            "ms_adapter": ("MS_ADAPTER",),
            "ms_embeds": ("MS_EMBEDS",),
            "scale": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 2.0, "step": 0.05}),
            "start_step": ("INT", {"default": 5, "min": 0, "max": 1024,
                                   "tooltip": "pseudo attention mask kicks in after this many "
                                              "denoising steps (only when mask_threshold > 0)"}),
        }}

    RETURN_TYPES = ("MODEL",)
    FUNCTION = "apply"
    CATEGORY = "MSdiffusion/Native"
    DESCRIPTION = ("Patch every native SDXL attn2 with MS-Diffusion dual-branch "
                   "cross-attention. Output is a normal MODEL: use stock KSampler / "
                   "ControlNet / FreeU / LoRA nodes downstream.")

    def apply(self, model, ms_adapter, ms_embeds, scale, start_step):
        import copy
        m = model.clone()
        desc = enumerate_attn2(m.model.diffusion_model)
        if not desc:
            raise RuntimeError("MS-Diffusion: no attn2 modules found. This adapter requires a stock SDXL UNet.")
        perm = build_block_weights(desc, ms_adapter.hidden_sizes)

        emb = copy.copy(ms_embeds)
        emb.scale = float(scale)
        emb.start_step = int(start_step)
        sid = new_state(ms_adapter, emb)

        # IMPORTANT: set_model_patch_replace() REPLACES the whole
        # transformer_options dict and silently drops previously attached keys.
        # Build patches_replace in place on the cloned options instead.
        to = m.model_options["transformer_options"]
        pr = dict(to.get("patches_replace", {}))
        a2 = dict(pr.get("attn2", {}))
        for pos, (key3, _inner, _section) in enumerate(desc):
            a2[key3] = make_attn2_patch(sid, perm[pos], is_sentinel=(pos == 0))
        pr["attn2"] = a2
        to["patches_replace"] = pr
        print("[MS-Diffusion] patched %d attn2 layers (scale=%.2f, layout=%s)"
              % (len(desc), emb.scale, "on" if emb.mask_threshold is not None else "off"))
        return (m,)


# No custom image-batch node: ComfyUI's native `ImageBatch` already stacks
# images into [N,H,W,3] and aligns size/channels. Chain it for N > 2.
# (The old `MS_Object_img_Batch` / `MSDiffusionImageBatch` ids are intentionally
#  gone; replace them with the native ImageBatch node in old graphs.)

NODE_CLASS_MAPPINGS = {
    "MSDiffusionLoader": MSDiffusionLoader,
    "MSDiffusionConditioning": MSDiffusionConditioning,
    "MSDiffusionApply": MSDiffusionApply,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "MSDiffusionLoader": "MS-Diffusion Loader (Native)",
    "MSDiffusionConditioning": "MS-Diffusion Conditioning (Native)",
    "MSDiffusionApply": "MS-Diffusion Apply to MODEL (Native)",
}