import logging
import os
import types

import comfy.ldm.modules.attention

log = logging.getLogger(__name__)


def _masked_attention(q, k, v, heads, mask, transformer_options={}, **kwargs):
    # Bypass wrap_attn (sage/etc may ignore masks) by calling attention_pytorch directly.
    return comfy.ldm.modules.attention.attention_pytorch(
        q, k, v, heads, mask=mask,
        _inside_attn_wrapper=True,
        transformer_options=transformer_options,
        **kwargs,
    )


def _wan_t2v_forward(self, mask_fn, x, context, transformer_options={}, **kwargs):
    q = self.norm_q(self.q(x))
    k = self.norm_k(self.k(context))
    v = self.v(context)

    mask = mask_fn(q.shape[1], k.shape[1], q.dtype, q.device, transformer_options)
    if mask is not None:
        x = _masked_attention(q, k, v, heads=self.num_heads, mask=mask,
                              transformer_options=transformer_options)
    else:
        x = comfy.ldm.modules.attention.optimized_attention(
            q, k, v, heads=self.num_heads, transformer_options=transformer_options,
        )
    return self.o(x)


def _wan_i2v_forward(self, mask_fn, x, context, context_img_len, transformer_options={}, **kwargs):
    context_img = context[:, :context_img_len]
    context_text = context[:, context_img_len:]

    q = self.norm_q(self.q(x))

    k_img = self.norm_k_img(self.k_img(context_img))
    v_img = self.v_img(context_img)
    img_x = comfy.ldm.modules.attention.optimized_attention(
        q, k_img, v_img, heads=self.num_heads, transformer_options=transformer_options,
    )

    k = self.norm_k(self.k(context_text))
    v = self.v(context_text)

    mask = mask_fn(q.shape[1], k.shape[1], q.dtype, q.device, transformer_options)
    if mask is not None:
        x = _masked_attention(q, k, v, heads=self.num_heads, mask=mask,
                              transformer_options=transformer_options)
    else:
        x = comfy.ldm.modules.attention.optimized_attention(
            q, k, v, heads=self.num_heads, transformer_options=transformer_options,
        )

    return self.o(x + img_x)


_MASK_CACHE_OFF = os.environ.get("CGLIDE_NO_MASK_CACHE", "") == "1"
_CACHE_MISS = object()


def _make_cached_mask_fn(mask_fn):
    """Memoise mask_fn results so repeated cross-attn calls skip its per-call work.

    NOTE: the result is NOT purely shape-dependent — mask_fn returns None on
    unconditional passes (via transformer_options["cond_or_uncond"]) and derives the
    token grid from transformer_options["grid_sizes"] on Wan. Both MUST be part of the
    key: cond and uncond passes share identical shapes when run separately, and caching
    across them would apply the relay mask to the negative pass (or drop it from the
    positive one).

    prompt_relay.create_mask_fn already caches the expensive penalty-matrix build
    internally, so this outer layer only skips the per-call branching/lookup — cheap,
    but it runs blocks x steps x cond/uncond times. Escape hatch: CGLIDE_NO_MASK_CACHE=1.
    """
    if _MASK_CACHE_OFF:
        return mask_fn

    cache = {}

    def _opts_key(opts):
        cou = opts.get("cond_or_uncond", None)
        try:
            cou = tuple(cou) if cou is not None else None
        except TypeError:
            cou = str(cou)
        gs = opts.get("grid_sizes", None)
        if gs is not None:
            try:
                gs = tuple(int(v) for v in gs)
            except Exception:
                gs = str(gs)
        return (cou, gs, opts.get("promptrelay_attn_type"))

    def cached(q_len, k_len, dtype, device, opts):
        key = (q_len, k_len, dtype, device, _opts_key(opts))
        val = cache.get(key, _CACHE_MISS)
        if val is _CACHE_MISS:
            val = mask_fn(q_len, k_len, dtype, device, opts)
            cache[key] = val  # None results cached too — skips recheck-and-discard.
        return val

    return cached


def _make_masked_override(prev_override):
    """transformer_options override that routes mask-bearing attention calls through
    attention_pytorch (sage/etc. drop arbitrary masks). Chains to a prior override
    when no mask is present so we don't clobber other backends."""
    def override(func, *args, **kwargs):
        if kwargs.get("mask") is not None:
            return comfy.ldm.modules.attention.attention_pytorch(*args, **kwargs)
        if prev_override is not None:
            return prev_override(func, *args, **kwargs)
        return func(*args, **kwargs)
    return override


# Debug tracing is OFF by default: writing a log file inside the attention hot path
# (every block x every step x cond/uncond) costs tens of thousands of disk writes per
# generation. Set CGLIDE_DEBUG=1 in the environment to re-enable when diagnosing.
_DEBUG = os.environ.get("CGLIDE_DEBUG", "") == "1"

if _DEBUG:
    def debug_log(msg):
        try:
            log_path = os.path.join(os.path.dirname(__file__), "debug_prompt_relay.log")
            with open(log_path, "a", encoding="utf-8") as f:
                f.write(msg + "\n")
        except Exception:
            pass
else:
    def debug_log(msg):
        pass


def _make_ltx_mask_wrapper(underlying, mask_fn, attr):
    """Wrap an existing LTX cross-attn forward (the default `CrossAttention.forward`
    or another node's patch — e.g. KJNodes NAG), injecting PromptRelay's additive
    mask via the `mask` kwarg the upstream signature already accepts.

    `underlying` must already be bound to its module — callable as
    `underlying(x, context=..., mask=..., ...)`.
    """
    def wrapped(_self, x, context=None, mask=None, pe=None, k_pe=None, transformer_options={}):
        if _DEBUG:
            debug_log(f"wrapped called: x.shape={list(x.shape)} context.shape={list(context.shape) if context is not None else None} mask_is_none={mask is None}")
        if context is not None:
            opts = {**transformer_options, "promptrelay_attn_type": attr}
            pr_mask = mask_fn(x.shape[1], context.shape[1], x.dtype, x.device, opts)
            if _DEBUG:
                debug_log(f"mask_fn returned pr_mask_is_none={pr_mask is None}")
            if pr_mask is not None:
                if _DEBUG:
                    debug_log(f"pr_mask info: shape={list(pr_mask.shape)} min={pr_mask.min().item()} max={pr_mask.max().item()} sum={pr_mask.sum().item()}")
                mask = pr_mask if mask is None else mask + pr_mask

        if mask is not None:
            prev = transformer_options.get("optimized_attention_override")
            transformer_options = {
                **transformer_options,
                "optimized_attention_override": _make_masked_override(prev),
            }

        return underlying(
            x, context=context, mask=mask, pe=pe, k_pe=k_pe,
            transformer_options=transformer_options,
        )

    wrapped._promptrelay_wrapper = True
    return wrapped


class _CrossAttnPatch:
    """Descriptor that binds (impl, mask_fn) as a method onto a Wan cross-attn module."""

    def __init__(self, impl, mask_fn):
        self.impl = impl
        self.mask_fn = mask_fn

    def __get__(self, obj, objtype=None):
        impl, mask_fn = self.impl, self.mask_fn

        def wrapped(self_module, *args, **kwargs):
            return impl(self_module, mask_fn, *args, **kwargs)

        return types.MethodType(wrapped, obj)


def detect_model_type(model):
    """Return (arch, patch_size, temporal_stride) for latent geometry.

    temporal_stride is the VAE's pixel→latent temporal compression factor,
    used to convert user-facing pixel frame counts to latent frames.
    """
    diff_model = model.model.diffusion_model

    if hasattr(diff_model, "patch_size") and not hasattr(diff_model, "patchifier"):
        return "wan", tuple(diff_model.patch_size), 4

    if hasattr(diff_model, "patchifier"):
        return "ltx", (1, 1, 1), int(diff_model.vae_scale_factors[0])

    raise ValueError(
        f"Unsupported model type: {type(diff_model).__name__}. "
        f"Currently supports Wan and LTX models."
    )


def _check_unpatched(model_clone, key):
    if key in getattr(model_clone, "object_patches", {}):
        raise RuntimeError(
            f"PromptRelay: cross-attention forward at '{key}' is already patched by "
            "another node. Stacking is not supported for this architecture — remove "
            "the conflicting node."
        )


def apply_patches(model_clone, arch, mask_fn):
    diffusion_model = model_clone.get_model_object("diffusion_model")

    # One shared cache for the whole model: every block requests the same mask shape,
    # so the mask is built exactly once per generation instead of blocks x steps x 2 times.
    mask_fn = _make_cached_mask_fn(mask_fn)

    if arch == "wan":
        from comfy.ldm.wan.model import WanI2VCrossAttention
        for idx, block in enumerate(diffusion_model.blocks):
            key = f"diffusion_model.blocks.{idx}.cross_attn.forward"
            _check_unpatched(model_clone, key)
            cross_attn = block.cross_attn
            impl = _wan_i2v_forward if isinstance(cross_attn, WanI2VCrossAttention) else _wan_t2v_forward
            model_clone.add_object_patch(key, _CrossAttnPatch(impl, mask_fn).__get__(cross_attn, cross_attn.__class__))
        return

    if arch == "ltx":
        to = model_clone.model_options["transformer_options"]
        to["promptrelay_mask_fn"] = mask_fn

        for idx, block in enumerate(diffusion_model.transformer_blocks):
            for attr in ("attn2", "audio_attn2"):
                module = getattr(block, attr, None)
                if module is None:
                    continue
                key = f"diffusion_model.transformer_blocks.{idx}.{attr}.forward"
                # get_model_object returns the prior patch if present, else the default bound forward.
                underlying = model_clone.get_model_object(key)
                wrapper = _make_ltx_mask_wrapper(underlying, mask_fn, attr)
                model_clone.add_object_patch(key, types.MethodType(wrapper, module))
        return

    raise ValueError(f"Unknown model arch: {arch}")