"""
Muse Collective — model-architecture attention patching for MusePromptRelay-style
segment masking. Independently written for MuseDirectorSamplerV10Final: detects whether
a loaded diffusion model is a Wan or LTX transformer and installs a cross-attention
wrapper on every transformer block so a caller-supplied mask function can bias
attention scores per call, without touching the base model's own weights.
"""
import logging
import types

import comfy.ldm.modules.attention

log = logging.getLogger(__name__)


def detect_model_type(model):
    """Inspect model.model.diffusion_model and return (arch_name, patch_size, vae_scale)."""
    dm = model.model.diffusion_model
    is_ltx = hasattr(dm, "patchifier")
    is_wan = hasattr(dm, "patch_size") and not is_ltx
    if is_wan:
        return "wan", tuple(dm.patch_size), 4
    if is_ltx:
        return "ltx", (1, 1, 1), int(dm.vae_scale_factors[0])
    raise ValueError(
        f"detect_model_type: unrecognized diffusion_model class '{type(dm).__name__}' "
        f"— only Wan and LTX transformers are supported."
    )


def _biased_attention(q, k, v, num_heads, bias, transformer_options):
    """Run attention with an additive score bias — routes through ComfyUI's own
    non-flash attention path since fused/flash kernels can't accept an arbitrary bias."""
    return comfy.ldm.modules.attention.attention_pytorch(
        q, k, v, num_heads, mask=bias,
        _inside_attn_wrapper=True,
        transformer_options=transformer_options,
    )


def _run_cross_attention(q, k, v, num_heads, bias, transformer_options):
    if bias is None:
        return comfy.ldm.modules.attention.optimized_attention(
            q, k, v, heads=num_heads, transformer_options=transformer_options,
        )
    return _biased_attention(q, k, v, num_heads, bias, transformer_options)


def _wan_forward_t2v(block, mask_fn, x, context, transformer_options={}, **_ignored):
    q = block.norm_q(block.q(x))
    k = block.norm_k(block.k(context))
    v = block.v(context)
    bias = mask_fn(q.shape[1], k.shape[1], q.dtype, q.device, transformer_options)
    out = _run_cross_attention(q, k, v, block.num_heads, bias, transformer_options)
    return block.o(out)


def _wan_forward_i2v(block, mask_fn, x, context, context_img_len, transformer_options={}, **_ignored):
    img_ctx, text_ctx = context[:, :context_img_len], context[:, context_img_len:]
    q = block.norm_q(block.q(x))

    k_img = block.norm_k_img(block.k_img(img_ctx))
    v_img = block.v_img(img_ctx)
    img_branch = comfy.ldm.modules.attention.optimized_attention(
        q, k_img, v_img, heads=block.num_heads, transformer_options=transformer_options,
    )

    k_text = block.norm_k(block.k(text_ctx))
    v_text = block.v(text_ctx)
    bias = mask_fn(q.shape[1], k_text.shape[1], q.dtype, q.device, transformer_options)
    text_branch = _run_cross_attention(q, k_text, v_text, block.num_heads, bias, transformer_options)

    return block.o(text_branch + img_branch)


def _bind_wan_forward(block, mask_fn, is_i2v):
    impl = _wan_forward_i2v if is_i2v else _wan_forward_t2v

    def bound(_self, *args, **kwargs):
        return impl(block, mask_fn, *args, **kwargs)

    return types.MethodType(bound, block)


def _bind_ltx_forward(original_forward, mask_fn, attn_slot):
    def bound(_self, x, context=None, mask=None, pe=None, k_pe=None, transformer_options={}):
        if context is not None:
            relay_opts = {**transformer_options, "promptrelay_attn_type": attn_slot}
            relay_bias = mask_fn(x.shape[1], context.shape[1], x.dtype, x.device, relay_opts)
            if relay_bias is not None:
                mask = relay_bias if mask is None else mask + relay_bias
        if mask is not None:
            existing_override = transformer_options.get("optimized_attention_override")

            def forced_masked(func, *a, **kw):
                if kw.get("mask") is not None:
                    return comfy.ldm.modules.attention.attention_pytorch(*a, **kw)
                if existing_override is not None:
                    return existing_override(func, *a, **kw)
                return func(*a, **kw)

            transformer_options = {**transformer_options, "optimized_attention_override": forced_masked}
        return original_forward(
            x, context=context, mask=mask, pe=pe, k_pe=k_pe,
            transformer_options=transformer_options,
        )

    return bound


def apply_patches(model_clone, arch, mask_fn):
    """Install `mask_fn` as an attention-score bias source on every cross-attention
    block of the given architecture. Mutates model_clone via ComfyUI's object-patch
    API (no in-place weight edits)."""
    diffusion_model = model_clone.get_model_object("diffusion_model")

    if arch == "wan":
        from comfy.ldm.wan.model import WanI2VCrossAttention
        for idx, block in enumerate(diffusion_model.blocks):
            patch_key = f"diffusion_model.blocks.{idx}.cross_attn.forward"
            is_i2v = isinstance(block.cross_attn, WanI2VCrossAttention)
            model_clone.add_object_patch(patch_key, _bind_wan_forward(block.cross_attn, mask_fn, is_i2v))
        return

    if arch == "ltx":
        model_clone.model_options["transformer_options"]["promptrelay_mask_fn"] = mask_fn
        for idx, block in enumerate(diffusion_model.transformer_blocks):
            for attn_slot in ("attn2", "audio_attn2"):
                module = getattr(block, attn_slot, None)
                if module is None:
                    continue
                patch_key = f"diffusion_model.transformer_blocks.{idx}.{attn_slot}.forward"
                original_forward = model_clone.get_model_object(patch_key)
                bound = _bind_ltx_forward(original_forward, mask_fn, attn_slot)
                model_clone.add_object_patch(patch_key, types.MethodType(bound, module))
        return

    raise ValueError(f"apply_patches: unsupported arch '{arch}' (expected 'wan' or 'ltx')")