"""Few-step MOSS-Image adapters (rapid-moss DMAD 4-step, lodestones 8-step DMD) for ComfyUI.

Works on native ComfyUI MOSS-Image models and the OpenMOSS HF loader.
An adapter is a folder in models/rapid_moss/ holding lora.safetensors, time_embedding.safetensors and
dataspace.safetensors. The node returns a patched MODEL plus the data-space SAMPLER and SIGMAS for
SamplerCustomAdvanced (use BasicGuider: no CFG, no negative prompt).
"""
import inspect
import math
import os

import comfy.patcher_extension
import comfy.samplers
import folder_paths
import torch
from safetensors.torch import load_file
from torch.nn import functional as F

FOLDER = 'rapid_moss'
folder_paths.add_model_folder_path(FOLDER, os.path.join(folder_paths.models_dir, FOLDER))
FILES = ('lora.safetensors', 'time_embedding.safetensors', 'dataspace.safetensors')
PROJ = ('q', 'k', 'v', 'o', 'ffn_up_gate', 'ffn_down')
TIME = ('time_embedding.0.weight', 'time_embedding.0.bias', 'time_embedding.2.weight', 'time_embedding.2.bias',
        'shared_block_modulation.bias', 'video_head_modulation.bias')
DM = 'diffusion_model.'
WRAPPER_KEY = 'rapid_moss'


def _adapter_dirs():
    found = {}
    for root in folder_paths.get_folder_paths(FOLDER):
        if os.path.isdir(root):
            for name in sorted(os.listdir(root)):
                path = os.path.join(root, name)
                if name not in found and all(os.path.isfile(os.path.join(path, f)) for f in FILES):
                    found[name] = path
    return found


class _DeviceCopy:
    """Tensors kept on the CPU, copied once to the device that first needs them."""

    def __init__(self, *tensors):
        self.cpu = tensors
        self.device = None
        self.copies = None

    def get(self, device):
        if device != self.device:
            self.copies = [t.to(device) for t in self.cpu]
            self.device = device
        return self.copies


# ---- OpenMOSS HF loader: patch the modules of the diffusers transformer ----

class _FP32Linear(torch.nn.Module):
    """fp32 linear outside ComfyUI's weight management. `weight` is a plain attribute so the MOSS
    forward (which casts its input to `.weight.dtype`) runs this path in fp32."""

    def __init__(self, weight, bias):
        super().__init__()
        self.weight = weight
        self.tensors = _DeviceCopy(weight, bias)

    def forward(self, x):
        weight, bias = self.tensors.get(x.device)
        return F.linear(x.float(), weight, bias)


def _lora_forward(module, down, up, scale):
    """Unmerged fp32 LoRA: the delta is below bf16 rounding, so it must not be merged into the weight."""
    class_forward = type(module).forward
    tensors = _DeviceCopy(down, up)

    def forward(x, *args, **kwargs):
        out = class_forward(module, x, *args, **kwargs)
        d, u = tensors.get(x.device)
        return out + (((x.float() @ d.t()) @ u.t()) * scale).to(out.dtype)
    return forward


def _patch_hf(patched, lora, time_path, strength):
    patched.add_object_patch(DM + 'time_embedding', torch.nn.Sequential(
        _FP32Linear(time_path['time_embedding.0.weight'].float(), time_path['time_embedding.0.bias'].float()),
        torch.nn.SiLU(),
        _FP32Linear(time_path['time_embedding.2.weight'].float(), time_path['time_embedding.2.bias'].float())))
    for name in ('shared_block_modulation', 'video_head_modulation'):
        weight = patched.get_model_object(DM + name).weight.detach().to('cpu', torch.float32)
        patched.add_object_patch(DM + name, _FP32Linear(weight, time_path[name + '.bias'].float()))
    for i in range(len(patched.get_model_object(DM + 'layers'))):
        for proj in PROJ:
            target = f'{DM}layers.{i}.proj.{proj}'
            down = lora[target + '.lora_down.weight'].float()
            up = lora[target + '.lora_up.weight'].float()
            scale = strength * float(lora[target + '.alpha']) / down.shape[0]
            patched.add_object_patch(target + '.forward', _lora_forward(patched.get_model_object(target), down, up, scale))


# ---- native ComfyUI MOSS-Image: the whole forward as a DIFFUSION_MODEL wrapper ----

def _timestep_embedding(timestep, dim=256):
    """The native sinusoidal embedding (timestep * 1000, cos | sin) in float64, never rounded to bf16."""
    half = dim // 2
    freqs = torch.exp(-math.log(10000) * torch.arange(half, dtype=torch.float64, device=timestep.device) / half)
    phase = (timestep.double() * 1000.0)[:, None] * freqs
    return torch.cat((phase.cos(), phase.sin()), dim=-1).float()


def _fp32_linear(module, x, bias):
    """The base modulation weight in fp32 with the adapter's bias (ComfyUI casts or loads the weight)."""
    import comfy.ops
    weight, base_bias, stream = comfy.ops.cast_bias_weight(module, x, dtype=torch.float32, offloadable=True)
    out = F.linear(x, weight, bias)
    comfy.ops.uncast_bias_weight(module, weight, base_bias, stream)
    return out


def _lin(module, x, lora):
    out = module(x)
    down, up = lora
    return out + F.linear(F.linear(x.to(down.dtype), down), up).to(out.dtype)


class _NativeForward:
    """MOSS-Image forward with the adapter: fp32 time path, modulation and residual stream, unmerged LoRA
    (as rapid-moss trains and samples it). ComfyUI's memory compiler (comfy-aimdo) records the allocations of
    each block and replays them; a tensor allocated in one block and freed in a later one aborts it
    ("aimdo memory compile error"). So, like the native blocks, every block updates a residual allocated before
    the compiled region in place, and the adapter tensors reach the device before that region."""

    def __init__(self, lora, time_path, strength, num_layers, lora_dtype=torch.bfloat16):
        self.time = _DeviceCopy(*(time_path[k].float() for k in TIME))
        self.lora = []
        for i in range(num_layers):
            tensors = []
            for proj in PROJ:
                prefix = f'{DM}layers.{i}.proj.{proj}'
                down = lora[prefix + '.lora_down.weight']
                scale = strength * float(lora[prefix + '.alpha']) / down.shape[0]
                tensors += [down.to(lora_dtype), (lora[prefix + '.lora_up.weight'].float() * scale).to(lora_dtype)]
            self.lora.append(_DeviceCopy(*tensors))

    def __call__(self, executor, x, timestep, context, transformer_options={}, **kwargs):
        import comfy.model_prefetch
        from comfy.ldm.moss_image.model import modulated_rms_norm

        dm = executor.class_obj
        B, _, H, W = x.shape
        dtype, device = x.dtype, x.device
        transformer_options = transformer_options.copy()
        if 'transformer_options' in inspect.signature(dm.rope).parameters:
            pe = dm.rope(H, W, context.shape[1], device, transformer_options)
        else:
            pe = dm.rope(H, W, context.shape[1], device)

        t0_w, t0_b, t2_w, t2_b, mod_bias, head_bias = self.time.get(device)
        lora = [layer.get(device) for layer in self.lora]
        vec = F.linear(F.silu(F.linear(_timestep_embedding(timestep), t0_w, t0_b)), t2_w, t2_b)
        mod = _fp32_linear(dm.shared_block_modulation, vec, mod_bias).unflatten(-1, (6, -1)).unsqueeze(2).unbind(1)
        head_shift, head_scale = _fp32_linear(dm.video_head_modulation, vec, head_bias).unflatten(-1, (2, -1)).unsqueeze(2).unbind(1)

        txt = dm.text_embedder(context, pe[:, H * W:], transformer_options)
        # fp32 residual, allocated outside the memory compiler region; the blocks update it in place
        x = torch.cat((dm.video_embedder(x.flatten(2).transpose(1, 2)), txt), dim=1).float()
        del txt

        blocks_replace = transformer_options.get('patches_replace', {}).get('dit', {})
        transformer_options['total_blocks'] = len(dm.layers)
        transformer_options['block_type'] = 'single'
        prefetch_queue = comfy.model_prefetch.make_prefetch_queue(list(dm.layers), device, transformer_options)
        comfy.model_prefetch.malloc_graph_begin(device)
        for i, block in enumerate(dm.layers):
            comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, block, dtype, malloc_scope='block')
            transformer_options['block_index'] = i
            if ('single_block', i) in blocks_replace:
                def block_wrap(args, block=block, layer=lora[i]):
                    return {'img': self.block(block, layer, args['img'], args['mod'], args['pe'], dtype, args['transformer_options'])}
                x = blocks_replace[('single_block', i)]({'img': x, 'vec': vec, 'mod': mod, 'pe': pe, 'transformer_options': transformer_options},
                                                        {'original_block': block_wrap})['img']
            else:
                self.block(block, lora[i], x, mod, pe, dtype, transformer_options)
        comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, None, malloc_scope='block')
        comfy.model_prefetch.malloc_graph_end()

        y = modulated_rms_norm(x[:, :H * W], dm.final_norm_video, head_scale, head_shift).to(dtype)
        return dm.final_linear_video(y).transpose(1, 2).reshape(B, dm.latent_channels, H, W)

    @staticmethod
    def block(block, lora, x, mod, pe, dtype, transformer_options):
        from comfy.ldm.moss_image.model import attention, modulated_rms_norm, situ

        q_lora, k_lora, v_lora, o_lora, up_lora, down_lora = zip(lora[0::2], lora[1::2])
        shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = mod
        y = modulated_rms_norm(x, block.norm1, scale_a, shift_a).to(dtype)
        if block.fused_qkv:
            q, k, v = block.proj['qkv'](y).chunk(3, dim=-1)
        else:
            q, k, v = block.proj['q'](y), block.proj['k'](y), block.proj['v'](y)
        y_l = y.to(q_lora[0].dtype)
        q, k, v = (out + F.linear(F.linear(y_l, down), up).to(out.dtype) for out, (down, up) in zip((q, k, v), (q_lora, k_lora, v_lora)))
        out = attention(q, k, v, block.q_norm, block.k_norm, pe, block.num_heads, block.comfy_attention, transformer_options)
        del q, k, v, y_l
        x.addcmul_(_lin(block.proj['o'], out, o_lora), gate_a)
        del out
        y = modulated_rms_norm(x, block.norm2, scale_f, shift_f).to(dtype)
        hidden, gate = _lin(block.proj['ffn_up_gate'], y, up_lora).chunk(2, dim=-1)
        x.addcmul_(_lin(block.proj['ffn_down'], situ(hidden, gate), down_lora), gate_f)
        return x


def _is_native(patched):
    try:
        from comfy.ldm.moss_image.model import MOSSImage
    except ImportError:
        return False
    return isinstance(patched.get_model_object(DM[:-1]), MOSSImage)


@torch.no_grad()
def sample_dataspace(model, x, sigmas, extra_args=None, callback=None, disable=None, Q=None, codes=None):
    """x0 = 0; for r: x = z + Q x0, v = net(x, code_r), x0 = z - v (one noise z, no schedule, no CFG)."""
    extra_args = {} if extra_args is None else extra_args
    z = x.float()  # sigmas[0] == 1, so the sampler receives pure noise
    mix = Q.to(z.device, torch.float32)
    x0 = torch.zeros_like(z)
    s_in = z.new_ones([z.shape[0]])
    for r in range(len(sigmas) - 1):
        x_in = z + torch.einsum('oc,bchw->bohw', mix, x0)
        s = float(codes[r]) / 1000.0  # flow model: timestep = sigma * 1000 = code
        denoised = model(x_in, s * s_in, **extra_args).float()  # = x_in - s * v
        x0 = z - (x_in - denoised) / s
        if callback is not None:
            callback({'x': x_in, 'i': r, 'sigma': sigmas[r], 'sigma_hat': sigmas[r], 'denoised': x0})
    return x0


class RapidMOSSAdapter:
    @classmethod
    def INPUT_TYPES(cls):
        names = list(_adapter_dirs()) or ['(no adapter folder in models/rapid_moss)']
        return {'required': {
            'model': ('MODEL',),
            'adapter': (names,),
            'steps': ('INT', {'default': 4, 'min': 1, 'max': 8,
                              'tooltip': 'Network evaluations; up to the adapter step-code count (4 for rapid-moss, 8 for lodestones).'}),
            'lora_strength': ('FLOAT', {'default': 1.0, 'min': 0.0, 'max': 2.0, 'step': 0.05}),
        }}

    RETURN_TYPES = ('MODEL', 'SAMPLER', 'SIGMAS')
    RETURN_NAMES = ('model', 'sampler', 'sigmas')
    FUNCTION = 'apply'
    CATEGORY = 'MOSS-Image/rapid-moss'
    DESCRIPTION = ('Few-step MOSS-Image: patches the MOSS-Image MODEL with a data-space adapter (unmerged LoRA, '
                   'fp32 time path) and returns the matching SAMPLER and SIGMAS for SamplerCustomAdvanced + BasicGuider.')

    def apply(self, model, adapter, steps, lora_strength):
        dirs = _adapter_dirs()
        if adapter not in dirs:
            raise FileNotFoundError(f'Put {", ".join(FILES)} in a folder under models/{FOLDER}/ (got {adapter!r}).')
        path = dirs[adapter]
        lora = load_file(os.path.join(path, 'lora.safetensors'))
        time_path = load_file(os.path.join(path, 'time_embedding.safetensors'))
        data = load_file(os.path.join(path, 'dataspace.safetensors'))
        codes = data['step_codes'].float()
        if steps > len(codes):
            raise ValueError(f'{adapter} has {len(codes)} step codes; steps must be <= {len(codes)}.')

        patched = model.clone()
        base = patched.get_model_object(DM + 'shared_block_modulation')
        if base.weight.dtype not in (torch.bfloat16, torch.float16, torch.float32):
            raise ValueError('Load MOSS-Image with weight_dtype default/bf16/fp32; fp8 weights are not supported here.')
        if _is_native(patched):
            patched.remove_wrappers_with_key(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, WRAPPER_KEY)
            patched.add_wrapper_with_key(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, WRAPPER_KEY, _NativeForward(
                lora, time_path, lora_strength, len(patched.get_model_object(DM + 'layers'))))
        else:
            _patch_hf(patched, lora, time_path, lora_strength)

        sampler = comfy.samplers.KSAMPLER(sample_dataspace, extra_options={'Q': data['latent_mix.Q'].float(), 'codes': codes})
        sigmas = torch.cat((codes[:steps] / 1000.0, torch.zeros(1)))
        return (patched, sampler, sigmas)


NODE_CLASS_MAPPINGS = {'RapidMOSSAdapter': RapidMOSSAdapter}
NODE_DISPLAY_NAME_MAPPINGS = {'RapidMOSSAdapter': 'MOSS-Image Few-Step Adapter (rapid-moss)'}