"""Avoid rebuilding immutable prefix layout during cached image decode."""
import math
import types
import torch
from diffusers.models.modeling_outputs import Transformer2DModelOutput
from diffusers.models.transformers.transformer_qwenimage21 import QwenImage21KVCache


def metadata(model, kw):
    cache = kw['kv_cache']
    if not hasattr(cache, '_fast_layout'):
        mask = kw['img_mask']
        repeats = torch.where(mask, 4, 1)[0]
        image_pad_mask = torch.repeat_interleave(mask[0], repeats)
        shapes = kw['img_shapes']
        rope = model.pos_embed(shapes[0], image_pad_mask, device=mask.device)
        _, target_mask = model.build_token_metadata(image_pad_mask, shapes[0])
        target_tokens = math.prod(shapes[0][-1])
        prefix = rope.shape[0] - target_tokens
        valid = None
        enc_mask = kw.get('encoder_hidden_states_mask')
        if enc_mask is not None:
            valid = torch.ones(kw['hidden_states'].shape[0], image_pad_mask.shape[0], dtype=torch.bool, device=mask.device)
            valid[:, (~image_pad_mask).nonzero(as_tuple=True)[0]] = enc_mask.bool()[:, ~mask[0][:enc_mask.shape[1]]]
            valid = valid[:, None, None, :]
        cache._fast_layout = (rope[prefix:], target_mask[prefix:], valid, target_tokens)
    return cache._fast_layout


def core(model, hidden_states, timestep, cache, layout):
    rope, target_mask, attention_mask, target_tokens = layout
    # Keep the original img_in shape, including condition-image tokens for edits.
    image = model.img_in(hidden_states)
    hidden = image[:, -target_tokens:]
    timestep = timestep.to(image.dtype)
    timestep = torch.cat([timestep, timestep.new_zeros(1)], dim=0)
    temb = model.time_text_embed(timestep, image)
    modulation = model.modulation(temb)
    for i, block in enumerate(model.transformer_blocks):
        hidden = block(hidden_states=hidden, modulation=modulation, rotary_emb=rope,
                       attention_mask=attention_mask, target_token_mask=target_mask,
                       layer_cache=cache.get_layer(i), kv_cache_mode='cached',
                       cache_write_slice=None, segments=None, key_valid=None)
    hidden = model.norm_out(hidden, temb, target_mask)
    return model.proj_out(hidden)


class Graph:
    def __init__(self, model, hidden, timestep, cache, layout):
        self.hidden = hidden.clone()
        self.timestep = timestep.clone()
        self.cache = QwenImage21KVCache(len(cache.layer_caches))
        for dst, src in zip(self.cache.layer_caches, cache.layer_caches):
            dst.store(src.k.clone(), src.v.clone())
        self.layout = tuple(x.clone() if isinstance(x, torch.Tensor) else x for x in layout)
        stream = torch.cuda.Stream()
        stream.wait_stream(torch.cuda.current_stream())
        with torch.cuda.stream(stream):
            for _ in range(2):
                core(model, self.hidden, self.timestep, self.cache, self.layout)
        torch.cuda.current_stream().wait_stream(stream)
        self.graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(self.graph, stream=stream):
            self.output = core(model, self.hidden, self.timestep, self.cache, self.layout)
        self.owner = cache

    def run(self, hidden, timestep, cache, layout):
        self.hidden.copy_(hidden)
        self.timestep.copy_(timestep)
        if self.owner is not cache:
            for dst, src in zip(self.cache.layer_caches, cache.layer_caches):
                dst.k.copy_(src.k)
                dst.v.copy_(src.v)
            for dst, src in zip(self.layout, layout):
                if isinstance(dst, torch.Tensor):
                    dst.copy_(src)
            self.owner = cache
        self.graph.replay()
        return self.output


def dispatch(self, **kw):
    if kw.get('kv_cache_mode') != 'cached' or kw.get('attention_kwargs'):
        return self._original_forward(**kw)
    cache = kw['kv_cache']
    layout = metadata(self, kw)
    hidden, timestep = kw['hidden_states'], kw['timestep']
    if self._fast_use_graph:
        key = (tuple(hidden.shape), tuple(cache.get_layer(0).k.shape),
               tuple(kw['img_shapes'][0]), layout[2] is not None)
        if key not in self._fast_graphs:
            if len(self._fast_graphs) >= 2:
                self._fast_graphs.pop(next(iter(self._fast_graphs)))
            self._fast_graphs[key] = Graph(self, hidden, timestep, cache, layout)
        result = self._fast_graphs[key].run(hidden, timestep, cache, layout)
    else:
        result = core(self, hidden, timestep, cache, layout)
    return (result,) if not kw.get('return_dict', True) else Transformer2DModelOutput(sample=result)


def enable(pipe, graph=False):
    m = pipe.transformer
    if not hasattr(m, '_original_forward'):
        m._original_forward = m.forward
        m.forward = types.MethodType(dispatch, m)
        m._fast_graphs = {}
    m._fast_use_graph = graph
    return pipe