"""
key_mapper.py — Shared Checkpoint ↔ LoRA Key Mapping

Single authoritative source for ALL key mapping logic, used by both the
extractor (:mod:`lora_extractor`) and the baker (:mod:`baking_processor_matching`).

Design follows ``plans/working_principles.md``:
    P1 (Delete First)  — replaces two separate implementations with one shared module
    P3 (One Pattern)   — single ``_ARCH_KEY_RULES`` registry derived from both directions
    P4 (Resist Guards) — clean prefix-matching logic, no special-case cascades
    P5 (Simple Dispatch) — each function has one clear responsibility

Architecture support:
    - SDXL UNet (model.diffusion_model.input_blocks.*)
    - SDXL TE1 (cond_stage_model.transformer.text_model.*)
    - SDXL TE2 (conditioner.embedders.*.transformer.text_model.*)
    - SD1.5 UNet (diffusion_model.input_blocks.*)
    - SD1.5 TE (cond_stage_model.transformer.text_model.*)
    - Flux standard (model.diffusion_model.transformer.double_blocks.* / single_blocks.*)
    - Flux Klein (bare double_blocks.* / single_blocks.* — adds transformer. prefix)
    - Z-Image / Lumina2 (layers.* / cap_embedder.* etc.)
    - Anima (net.blocks.*)
"""

import re
from typing import Dict, List, Optional, Tuple


# ============================================================================
# Shared helpers
# ============================================================================

_TENSOR_SUFFIXES = ('.weight', '.bias', '.alpha')


def strip_tensor_suffix(key: str) -> str:
    """
    Strip ``.weight``, ``.bias``, or ``.alpha`` suffix from a tensor key.

    Shared helper used by key_mapper functions, the baker's shape-index
    builders, and the extractor's key matching.

    Example:
        >>> strip_tensor_suffix("double_blocks.0.img_attn.proj.weight")
        "double_blocks.0.img_attn.proj"
    """
    for suffix in _TENSOR_SUFFIXES:
        if key.endswith(suffix):
            return key[:-len(suffix)]
    return key


# ============================================================================
# Pattern matching helpers (for diffusers key conversion)
# ============================================================================


def _match_pattern(key: str, pattern: str) -> Optional[Dict[str, str]]:
    """
    Match a ``{i}`` pattern against a key.

    E.g., ``_match_pattern("double_blocks.3.img_attn.qkv", "double_blocks.{i}")``
    returns ``{"i": "3"}``.

    Returns:
        Dict of ``{placeholder: value}`` if matched, ``None`` otherwise.
    """
    escaped = re.escape(pattern).replace(r"\{i\}", r"(\d+)")
    m = re.match(escaped, key)
    if m:
        return {"i": m.group(1)}
    return None


def _fill_pattern(pattern: str, vars: Dict[str, str]) -> str:
    """Fill ``{i}`` placeholders in a pattern with extracted values."""
    result = pattern
    for k, v in vars.items():
        result = result.replace("{" + k + "}", v)
    return result


def _strip_known_prefixes(key: str, arch: str) -> str:
    """
    Strip known LoRA/architecture prefixes for diffusers key matching.

    For Flux: strips ``transformer.``, ``model.diffusion_model.``, etc.
    """
    rules = _ARCH_KEY_RULES.get(arch, {})
    lora_prefix = rules.get("lora_prefix", "")
    if key.startswith(lora_prefix):
        return key[len(lora_prefix):]
    # Also try bare model.diffusion_model. prefix
    if key.startswith("model.diffusion_model."):
        return key[len("model.diffusion_model."):]
    if key.startswith("diffusion_model."):
        return key[len("diffusion_model."):]
    return key


# ============================================================================
# Architecture key rules — SINGLE source of truth
# ============================================================================
# Each entry defines:
#   checkpoint_prefixes: List of (checkpoint_prefix, lora_prefix).
#     Forward mapping: if a key starts with checkpoint_prefix,
#     replace it with lora_prefix.
#   bare_paths: List of path prefixes that may appear WITHOUT any model
#     prefix (e.g. Klein Flux "double_blocks.*").  For forward mapping,
#     these get ``lora_prefix`` prepended.
#   lora_prefix: Canonical LoRA prefix for this architecture.
#     Used when a bare_path is matched.

_ARCH_KEY_RULES: Dict[str, Dict] = {
    # ── Flux (standard + Klein) ──────────────────────────────────────────
    # Standard: model.diffusion_model.transformer.double_blocks.*
    # Normalized (after normalizer strips model.diffusion_model.):
    #   transformer.double_blocks.*
    # Klein (bare, no prefix): double_blocks.*
    "flux": {
        "checkpoint_prefixes": [
            # Order matters: more specific first
            ("model.diffusion_model.transformer.double_blocks.", "transformer.double_blocks."),
            ("model.diffusion_model.transformer.single_blocks.", "transformer.single_blocks."),
            ("transformer.double_blocks.", "transformer.double_blocks."),
            ("transformer.single_blocks.", "transformer.single_blocks."),
        ],
        "bare_paths": [
            "double_blocks.",    # Klein Flux bare key → transformer.double_blocks.*
            "single_blocks.",    # Klein Flux bare key → transformer.single_blocks.*
        ],
        "lora_prefix": "transformer.",
        # ── Diffusers-format LoRA keys (for ComfyUI load_lora compatibility) ──
        "diffusers": {
            "block_maps": {
                "double_blocks.{i}": "transformer_blocks.{i}",
                "single_blocks.{i}": "single_transformer_blocks.{i}",
            },
            "qkv_splits": {
                "double_blocks.{i}.img_attn.qkv": [
                    ("transformer_blocks.{i}.attn.to_q", (0, 1)),
                    ("transformer_blocks.{i}.attn.to_k", (1, 1)),
                    ("transformer_blocks.{i}.attn.to_v", (2, 1)),
                ],
                "double_blocks.{i}.txt_attn.qkv": [
                    ("transformer_blocks.{i}.attn.add_q_proj", (0, 1)),
                    ("transformer_blocks.{i}.attn.add_k_proj", (1, 1)),
                    ("transformer_blocks.{i}.attn.add_v_proj", (2, 1)),
                ],
                "single_blocks.{i}.linear1": [
                    ("single_transformer_blocks.{i}.attn.to_q", (0, 1)),
                    ("single_transformer_blocks.{i}.attn.to_k", (1, 1)),
                    ("single_transformer_blocks.{i}.attn.to_v", (2, 1)),
                    ("single_transformer_blocks.{i}.proj_mlp", (3, 4)),
                ],
            },
            "simple_map": {
                # Double block sub-keys
                "img_attn.proj": "attn.to_out.0",
                "txt_attn.proj": "attn.to_add_out",
                "img_mod.lin": "norm1.linear",
                "txt_mod.lin": "norm1_context.linear",
                "img_mlp.0": "ff.net.0.proj",
                "img_mlp.2": "ff.net.2",
                "txt_mlp.0": "ff_context.net.0.proj",
                "txt_mlp.2": "ff_context.net.2",
                "img_attn.norm.query_norm": "attn.norm_q",
                "img_attn.norm.key_norm": "attn.norm_k",
                "txt_attn.norm.query_norm": "attn.norm_added_q",
                "txt_attn.norm.key_norm": "attn.norm_added_k",
                # Single block sub-keys
                "linear2": "proj_out",
                "modulation.lin": "norm.linear",
                "norm.query_norm": "attn.norm_q",
                "norm.key_norm": "attn.norm_k",
            },
            "basic_map": {
                # CRITICAL: The checkpoint stores non-block keys using DIFFUSERS
                # naming (e.g. "final_layer.linear", "img_in", "txt_in",
                # "time_in.in_layer"), NOT Klein checkpoint naming ("proj_out",
                # "x_embedder", "context_embedder").  We match the diffusers
                # name and map to the CHECKPOINT name because ComfyUI's
                # model_lora_keys_unet() builds LoRA key_map entries using the
                # CHECKPOINT key name (from flux_to_diffusers() MAP_BASIC tuples
                # where k[1]=checkpoint_name is the dict key).
                #
                # Resulting LoRA key = lora_prefix + new_basic = "transformer.{checkpoint_name}"
                # which matches ComfyUI's key_map["transformer.{checkpoint_name}"].
                "final_layer.linear": "proj_out",
                "img_in": "x_embedder",
                "txt_in": "context_embedder",
                "time_in.in_layer": "time_text_embed.timestep_embedder.linear_1",
                "time_in.out_layer": "time_text_embed.timestep_embedder.linear_2",
                "vector_in.in_layer": "time_text_embed.text_embedder.linear_1",
                "vector_in.out_layer": "time_text_embed.text_embedder.linear_2",
                "guidance_in.in_layer": "time_text_embed.guidance_embedder.linear_1",
                "guidance_in.out_layer": "time_text_embed.guidance_embedder.linear_2",
                "norm_out.linear": "norm_out.linear",
                # final_layer.adaLN_modulation.1 → checkpoint name norm_out.linear
                # (with swap_scale_shift in ComfyUI; we don't apply the transform
                #  so the LoRA tensor is the raw delta)
                "final_layer.adaLN_modulation.1": "norm_out.linear",
                # Flux 2 modulation keys (not in current ComfyUI MAP_BASIC,
                # but properly prefixed for future compatibility)
                "double_stream_modulation_img.lin": "double_stream_modulation_img.lin",
                "double_stream_modulation_txt.lin": "double_stream_modulation_txt.lin",
                "single_stream_modulation.lin": "single_stream_modulation.lin",
            },
        },
    },
    # ── Lumina2 ──────────────────────────────────────────────────────────
    # Normalized keys are bare (layers.*, cap_embedder.*, etc.) after
    # normalizer strips model.diffusion_model.  LoRA keys need the
    # diffusion_model. prefix to match ComfyUI's model parameters.
    "lumina2": {
        "checkpoint_prefixes": [
            ("layers.",           "diffusion_model.layers."),
            ("cap_embedder.",     "diffusion_model.cap_embedder."),
            ("context_refiner.",  "diffusion_model.context_refiner."),
            ("noise_refiner.",    "diffusion_model.noise_refiner."),
            ("final_layer.",      "diffusion_model.final_layer."),
            ("x_embedder",        "diffusion_model.x_embedder"),
            ("x_pad_token",       "diffusion_model.x_pad_token"),
            ("cap_pad_token",     "diffusion_model.cap_pad_token"),
        ],
        "bare_paths": [],
        "lora_prefix": "diffusion_model.",
        # ── Diffusers-format LoRA keys (for ComfyUI load_lora compatibility) ──
        "diffusers": {
            "block_maps": {
                # Lumina2 block prefix is pass-through (same in checkpoint and diffusers)
                "layers.{i}": "layers.{i}",
                "context_refiner.{i}": "context_refiner.{i}",
                "noise_refiner.{i}": "noise_refiner.{i}",
            },
            "qkv_splits": {
                "layers.{i}.attention.qkv": [
                    ("layers.{i}.attention.to_q", (0, 1)),
                    ("layers.{i}.attention.to_k", (1, 1)),
                    ("layers.{i}.attention.to_v", (2, 1)),
                ],
                "context_refiner.{i}.attention.qkv": [
                    ("context_refiner.{i}.attention.to_q", (0, 1)),
                    ("context_refiner.{i}.attention.to_k", (1, 1)),
                    ("context_refiner.{i}.attention.to_v", (2, 1)),
                ],
                "noise_refiner.{i}.attention.qkv": [
                    ("noise_refiner.{i}.attention.to_q", (0, 1)),
                    ("noise_refiner.{i}.attention.to_k", (1, 1)),
                    ("noise_refiner.{i}.attention.to_v", (2, 1)),
                ],
            },
            "simple_map": {
                # Checkpoint sub-key → Diffusers sub-key
                # Derived from z_image_to_diffusers() block_map (reversed)
                "attention.out": "attention.to_out.0",
                "attention.q_norm": "attention.norm_q",
                "attention.k_norm": "attention.norm_k",
                # Pass-through entries (identity in ComfyUI's block_map)
                "attention_norm1": "attention_norm1",
                "attention_norm2": "attention_norm2",
                "feed_forward.w1": "feed_forward.w1",
                "feed_forward.w2": "feed_forward.w2",
                "feed_forward.w3": "feed_forward.w3",
                "ffn_norm1": "ffn_norm1",
                "ffn_norm2": "ffn_norm2",
                "adaLN_modulation.0": "adaLN_modulation.0",
            },
            "basic_map": {
                # CRITICAL: Identity mapping — ComfyUI's model_lora_keys_unet()
                # indexes by the KEY of z_image_to_diffusers() key_map, which
                # is the CHECKPOINT key name (second element of MAP_BASIC tuples).
                # The diffusers rename happens in the VALUE, not the LoRA key.
                "all_final_layer.2-1.linear": "all_final_layer.2-1.linear",
                "all_final_layer.2-1.adaLN_modulation.1": "all_final_layer.2-1.adaLN_modulation.1",
                "all_x_embedder.2-1": "all_x_embedder.2-1",
                "cap_embedder.0": "cap_embedder.0",
                "cap_embedder.1": "cap_embedder.1",
                "x_pad_token": "x_pad_token",
                "cap_pad_token": "cap_pad_token",
                "t_embedder.mlp.0": "t_embedder.mlp.0",
                "t_embedder.mlp.2": "t_embedder.mlp.2",
            },
        },
    },
    # ── SDXL ─────────────────────────────────────────────────────────────
    # Checkpoint keys keep model.diffusion_model.* prefix.
    # LoRA keys use diffusion_model.* (strip model. prefix).
    "sdxl": {
        "checkpoint_prefixes": [
            ("model.diffusion_model.", "diffusion_model."),
        ],
        "bare_paths": [],
        "lora_prefix": "diffusion_model.",
    },
    # ── SD1.5 ────────────────────────────────────────────────────────────
    # SD1.5 UNet: diffusion_model.xxx → diffusion_model.xxx (pass-through)
    # SD1.5 CLIP: cond_stage_model.transformer.text_model.xxx
    #             → text_encoders.clip_l.transformer.text_model.xxx
    "sd15": {
        "checkpoint_prefixes": [
            ("cond_stage_model.transformer.text_model.",
             "text_encoders.clip_l.transformer.text_model."),
            ("diffusion_model.", "diffusion_model."),  # pass-through
        ],
        "bare_paths": [],
        "lora_prefix": "diffusion_model.",
    },
    # ── Anima ────────────────────────────────────────────────────────────
    # Anima: net.blocks.xxx → diffusion_model.blocks.xxx
    # After checkpoint_normalizer strips "net." prefix, keys are bare
    # blocks.xxx. The bare_paths rule catches these normalized keys.
    "anima": {
        "checkpoint_prefixes": [
            ("net.blocks.", "diffusion_model.blocks."),
        ],
        "bare_paths": [
            "blocks.",  # normalized bare blocks.* → diffusion_model.blocks.*
        ],
        "lora_prefix": "diffusion_model.",
    },
    # ── Additional TE/CLIP prefix patterns ───────────────────────────────
    # These are shared across architectures for text-encoder keys.
    # SDXL TE2: model.conditioner.embedders.0.transformer.text_model.xxx
    #           → text_encoders.transformer.text_model.xxx
    # SDXL TE1: model.conditioner.embedders.1.transformer.text_model.xxx
    #           → text_encoders.transformer.text_model.xxx
    # SDXL / SD1.5: model.text_model.xxx → text_model.xxx
    "te": {  # special — matches any arch's text encoder
        "checkpoint_prefixes": [
            ("model.conditioner.embedders.0.transformer.text_model.",
             "text_encoders.transformer.text_model."),
            ("model.conditioner.embedders.1.transformer.text_model.",
             "text_encoders.transformer.text_model."),
            ("model.conditioner.", "conditioner."),
            ("model.text_model.", "text_model."),
        ],
        "bare_paths": [],
        "lora_prefix": "",
    },
}

# Architectures in detection order (most specific first)
_ARCH_ORDER = ["flux", "lumina2", "anima", "sdxl", "sd15"]


# ============================================================================
# Forward mapping: checkpoint key → LoRA key  (used by extractor)
# ============================================================================

def checkpoint_key_to_lora_key(
    ckpt_key: str,
    arch: Optional[str] = None,
) -> str:
    """
    Convert a normalized checkpoint key to the equivalent LoRA master-format key.

    After :func:`~engine.checkpoint_normalizer.normalize_checkpoint_key`, the
    key reflects the canonical form of the *checkpoint* naming.  This function
    maps it to the *LoRA* master format used by :func:`~engine.identity_normalizer.identity_normalize`
    and the rest of the merger suite.

    Args:
        ckpt_key: Checkpoint key (may include ``.weight`` / ``.bias`` / ``.alpha`` suffix).
        arch: Optional architecture hint (``"flux"``, ``"sdxl"``, ``"sd15"``, etc.).
            If ``None``, tries all known architectures.

    Returns:
        LoRA master-format key (without ``.weight/.bias/.alpha`` suffix).
        If no known conversion applies, the key is returned unchanged
        (pass-through for already-LoRA-format keys).
    """
    # ── Strip tensor-type suffix ──────────────────────────────────────────
    key = strip_tensor_suffix(ckpt_key)

    # ── Apply architecture-specific prefix rules ─────────────────────────
    architectures_to_try: List[Tuple[str, Dict]] = []
    if arch is not None and arch in _ARCH_KEY_RULES:
        architectures_to_try = [(arch, _ARCH_KEY_RULES[arch])]
    else:
        # Try all architectures in order
        for a in _ARCH_ORDER:
            if a in _ARCH_KEY_RULES:
                architectures_to_try.append((a, _ARCH_KEY_RULES[a]))
        # Also try the generic 'te' rule
        if 'te' in _ARCH_KEY_RULES:
            architectures_to_try.append(('te', _ARCH_KEY_RULES['te']))

    for arch_name, rules in architectures_to_try:
        # ── Checkpoint prefix match ──────────────────────────────────────
        for ckpt_prefix, lora_prefix in rules.get("checkpoint_prefixes", []):
            if key.startswith(ckpt_prefix):
                return lora_prefix + key[len(ckpt_prefix):]

        # ── Bare path match (Klein-style, no prefix at all) ──────────────
        for bare_path in rules.get("bare_paths", []):
            if key.startswith(bare_path):
                lora_prefix = rules.get("lora_prefix", "")
                return lora_prefix + key

    # Pass-through: already in LoRA format (Flux bare, Z-Image, SD1.5,
    # clip_l/clip_g/t5xxl/text_encoder/text_encoders, etc.)
    return key


# ============================================================================
# Diffusers key conversion: checkpoint-internal → diffusers LoRA format
# (used by extractor for Flux/Lumina2 where ComfyUI expects diffusers keys)
# ============================================================================


def checkpoint_key_to_diffusers_key(
    ckpt_base: str,
    arch: str,
) -> List[Tuple[str, Optional[Tuple[int, int]]]]:
    """
    Convert a checkpoint-internal key base to diffusers-format LoRA key(s).

    For QKV-split keys (merged qkv → to_q/to_k/to_v), returns multiple entries
    with ``(offset_in_h, width_in_h)`` tuples. Non-split keys return
    ``(diff_key, None)``.

    The ``(offset_in_h, width_in_h)`` tuple means:
        row_start = offset_in_h * hidden_size
        row_end   = (offset_in_h + width_in_h) * hidden_size

    This tells the caller how to slice the SVD output along the output
    dimension (lora_B row dimension).

    Args:
        ckpt_base: Checkpoint-internal key base, **with** the ``lora_prefix``
            already applied (e.g. ``"transformer.double_blocks.0.img_attn.qkv"``).
        arch: Architecture name from ``_ARCH_KEY_RULES``.

    Returns:
        List of ``(diffusers_key, (offset_in_h, width_in_h) or None)``.
        If no diffusers rules apply, returns ``[(ckpt_base, None)]``
        (pass-through for architectures that don't need conversion).
    """
    rules = _ARCH_KEY_RULES.get(arch, {}).get("diffusers")
    if not rules:
        return [(ckpt_base, None)]

    lora_prefix = _ARCH_KEY_RULES[arch].get("lora_prefix", "")
    inner = _strip_known_prefixes(ckpt_base, arch)

    # ── 1. Check qkv_splits (1→N expansion) ──────────────────────────────
    for pattern, splits in rules.get("qkv_splits", {}).items():
        vars = _match_pattern(inner, pattern)
        if vars is not None:
            result = []
            for diff_pattern, (offset_h, width_h) in splits:
                diff_inner = _fill_pattern(diff_pattern, vars)
                result.append((f"{lora_prefix}{diff_inner}", (offset_h, width_h)))
            return result

    # ── 2. Check block_maps + simple_map (1→1 rename) ────────────────────
    for block_old, block_new in rules.get("block_maps", {}).items():
        vars = _match_pattern(inner, block_old)
        if vars is not None:
            block_old_filled = _fill_pattern(block_old, vars)
            inner_suffix = inner[len(block_old_filled):]  # e.g. ".img_attn.proj"

            for old_sub, new_sub in rules.get("simple_map", {}).items():
                if inner_suffix == f".{old_sub}" or inner_suffix.startswith(f".{old_sub}."):
                    suffix = inner_suffix[len(f".{old_sub}"):]
                    block_new_filled = _fill_pattern(block_new, vars)
                    diff_inner = f"{block_new_filled}.{new_sub}{suffix}"
                    return [(f"{lora_prefix}{diff_inner}", None)]

            # Block matched but no sub-match — pass through sub-key as-is
            block_new_filled = _fill_pattern(block_new, vars)
            return [(f"{lora_prefix}{block_new_filled}{inner_suffix}", None)]

    # ── 3. Check basic_map (non-block root-level keys) ───────────────────
    for old_basic, new_basic in rules.get("basic_map", {}).items():
        if inner == old_basic or inner.startswith(old_basic + "."):
            suffix = inner[len(old_basic):]
            result_key = f"{lora_prefix}{new_basic}{suffix}"
            return [(result_key, None)]

    # Pass-through: no diffusers rule matched
    return [(f"{lora_prefix}{inner}", None)]


# ============================================================================
# Reverse mapping: generate all LoRA key variants  (used by baker)
# ============================================================================

def build_lora_key_variants(ckpt_base: str) -> List[str]:
    """
    Generate ALL known LoRA key variants for a given checkpoint base key.

    Used by the baker to build its reverse key map.  Mirrors ComfyUI's
    ``model_lora_keys_unet()`` and ``model_lora_keys_clip()`` logic.

    Args:
        ckpt_base: Checkpoint base key (without ``.weight`` / ``.bias`` / ``.alpha`` suffix).

    Returns:
        List of LoRA key variants (canonical/master format first).
        Empty list if the key doesn't match any known architecture pattern.
    """
    variants: List[str] = []
    seen: set = set()

    def _add(v: str) -> None:
        if v not in seen:
            seen.add(v)
            variants.append(v)

    # ── Section 1: SDXL UNet (model.diffusion_model.*) ───────────────────
    if ckpt_base.startswith('model.diffusion_model.'):
        inner = ckpt_base[len('model.'):]  # diffusion_model.input_blocks.X
        pure_path = inner[len('diffusion_model.'):]  # input_blocks.X

        _add(inner)                                          # diffusion_model.*
        _add(pure_path)                                      # pure path (input_blocks.X)
        _add(f"lora_unet_{pure_path}")                       # Kohya-style

        # Underscore variants
        underscorized = pure_path.replace('.', '_')
        if underscorized != pure_path:
            _add(underscorized)
            _add(f"lora_unet_{underscorized}")

    # ── Section 2a: SD1.5 TE (cond_stage_model.transformer.*) ────────────
    if ckpt_base.startswith('cond_stage_model.transformer.'):
        inner = ckpt_base[len('cond_stage_model.'):]  # transformer.text_model.X
        pure_te = ckpt_base[len('cond_stage_model.transformer.'):]  # text_model.X

        _add(pure_te)                                        # text_model.*
        _add(inner)                                          # transformer.text_model.*
        _add(f"lora_te_{pure_te}")                           # Kohya-style
        _add(f"lora_te1_{pure_te}")                          # te1 variant
        _add(f"te.{pure_te}")                                # te. prefix
        _add(f"te1.{pure_te}")                               # te1. prefix

        # Underscore variants
        underscorized = pure_te.replace('.', '_')
        if underscorized != pure_te:
            _add(underscorized)
            _add(f"lora_te_{underscorized}")
            _add(f"lora_te1_{underscorized}")

    # ── Section 2a-ii: CLIP-L / CLIP-G (clip_l.text_model.* / clip_g.text_model.*) ──
    if ckpt_base.startswith(('clip_l.text_model.', 'clip_g.text_model.')):
        clip_prefix = ckpt_base.split('.')[0]  # clip_l or clip_g
        pure_clip = ckpt_base[len(f"{clip_prefix}."):]  # text_model.X

        _add(pure_clip)                                      # text_model.*
        _add(f"lora_te_{pure_clip}")                         # lora_te_ variant
        _add(f"te.{pure_clip}")                              # te. prefix
        _add(ckpt_base)                                      # original
        _add(f"lora_te_{clip_prefix}_{pure_clip}")           # lora_te_clip_l_ variant

        # Underscore variants
        underscorized = pure_clip.replace('.', '_')
        if underscorized != pure_clip:
            _add(underscorized)
            _add(f"lora_te_{underscorized}")

    # ── Section 2b: SDXL TE2 (conditioner.embedders.*) ───────────────────
    if ckpt_base.startswith('conditioner.embedders.'):
        parts = ckpt_base.split('.')
        if len(parts) >= 4 and parts[2].isdigit():
            embedder_idx = parts[2]
            inner_key = '.'.join(parts[3:])  # transformer.text_model.X
            pure_te2 = '.'.join(parts[4:])   # text_model.X

            _add(f"conditioner.embedders.{embedder_idx}.{inner_key}")
            _add(f"te2.{pure_te2}")
            _add(f"te1.{pure_te2}")                           # some use te1 for first embedder
            _add(f"lora_te2_{pure_te2}")
            _add(f"conditioner.{inner_key}")

    # ── Section 2c: T5 text encoder (t5xxl.transformer.*) ────────────────
    if ckpt_base.startswith('t5xxl.transformer.'):
        inner = ckpt_base[len('t5xxl.'):]  # transformer.encoder.layers.X
        pure_t5 = ckpt_base[len('t5xxl.transformer.'):]  # encoder.layers.X

        _add(inner)                                          # transformer.encoder.layers.X
        _add(pure_t5)                                        # encoder.layers.X
        _add(f"lora_te_{pure_t5}")                           # lora_te_ variant
        _add(f"te.{pure_t5}")                                # te. prefix

        # Underscore variants
        underscorized = pure_t5.replace('.', '_')
        if underscorized != pure_t5:
            _add(underscorized)
            _add(f"lora_te_{underscorized}")
            _add(f"te_{underscorized}")

    # ── Section 3: Flux (model.diffusion_model.transformer.*) ────────────
    if ckpt_base.startswith('model.diffusion_model.transformer.'):
        inner = ckpt_base[len('model.diffusion_model.'):]  # transformer.double_blocks.X
        pure_flux = ckpt_base[len('model.diffusion_model.transformer.'):]  # double_blocks.X

        _add(inner)                                          # transformer.double_blocks.X
        _add(pure_flux)                                      # double_blocks.X (bare)
        _add(f"lora_unet_{pure_flux}")                       # Kohya-style
        _add(f"lora_unet_transformer.{pure_flux}")            # rare Kohya variant
        _add(f"transformer.{pure_flux}")                     # explicit transformer. prefix

    # ── Section 4a: Z-Image / Musubi (model.diffusion_model.layers.*) ───
    if ckpt_base.startswith('model.diffusion_model.layers.'):
        inner = ckpt_base[len('model.diffusion_model.'):]  # layers.X
        pure_zi = ckpt_base[len('model.diffusion_model.layers.'):]  # X

        _add(inner)                                          # layers.*
        _add(pure_zi)                                        # pure path
        _add(f"lora_unet_{pure_zi}")                         # lora_unet variant
        _add(f"lora_unet_layers_{pure_zi}")                  # lora_unet_layers variant

    # ── Section 4b: Bare Z-Image / Musubi (layers.* without model prefix) ──
    if ckpt_base.startswith('layers.') and not ckpt_base.startswith('model.diffusion_model.layers.'):
        _add(ckpt_base)                                      # bare layers.*

    # ── Section 4c: Flux transformer entries (transformer.double_blocks.* / single_blocks.*) ──
    if ckpt_base.startswith('transformer.double_blocks.') or ckpt_base.startswith('transformer.single_blocks.'):
        _add(ckpt_base)                                      # bare transformer.*
        _add(f"model.diffusion_model.{ckpt_base}")           # with model.diffusion_model prefix
        _add(f"diffusion_model.{ckpt_base}")                 # with diffusion_model prefix

        pure_path = ckpt_base[len('transformer.'):]  # double_blocks.X
        _add(f"lora_unet_{pure_path}")                       # lora_unet variant

    # ── Section 4d: Flux keys with diffusion_model.{pure_path} (Klein normalized) ──
    if ckpt_base.startswith('transformer.double_blocks.') or ckpt_base.startswith('transformer.single_blocks.'):
        pure_path = ckpt_base[len('transformer.'):]  # double_blocks.X
        _add(f"diffusion_model.{pure_path}")                # diffusion_model.double_blocks.X
        _add(f"model.diffusion_model.{pure_path}")          # model.diffusion_model.double_blocks.X

    # ── Section 4e: Flux bare double_blocks/single_blocks (Klein 4B/9B format) ──
    if ckpt_base.startswith(('double_blocks.', 'single_blocks.')):
        _add(f"diffusion_model.{ckpt_base}")                # diffusion_model.double_blocks.X
        _add(f"model.diffusion_model.{ckpt_base}")          # model.diffusion_model.double_blocks.X
        _add(f"lora_unet_{ckpt_base}")                      # lora_unet_double_blocks.X

    # ── Section 5: Bare keys (generic fallback) ───────────────────────────
    # If the key doesn't have any standard prefix, register it as-is.
    if not ckpt_base.startswith((
        'model.', 'diffusion_model.', 'cond_stage_model.',
        'conditioner.', 'transformer.', 'layers.',
        'double_blocks.', 'single_blocks.',
    )):
        _add(ckpt_base)

    # ── Section 6: Anima block entries ────────────────────────────────────
    if ckpt_base.startswith('model.diffusion_model.'):
        inner = ckpt_base[len('model.diffusion_model.'):]
        if not inner.startswith(('input_blocks.', 'output_blocks.', 'middle_block.',
                                 'transformer.', 'layers.')):
            _add(f"lora_unet_{inner}")  # Anima-style block path

    # ── Additional: llm_adapter / llama entries ───────────────────────────
    if 'llm_adapter' in ckpt_base or 'llama' in ckpt_base:
        _add(ckpt_base)

    return variants