# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

import weakref
from types import SimpleNamespace
from unittest.mock import Mock

import pytest
import torch
from diffusers.models.autoencoders import AutoencoderKLWan

from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
    DistributedAutoencoderKLWan,
    OmniAutoencoderKLWan,
)
from vllm_omni.diffusion.models.interface import supports_chunked_vae_decode

pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]


@pytest.fixture
def vae():
    # Real causal convolutions and caches, with small random weights for CPU CI.
    return OmniAutoencoderKLWan(
        base_dim=4,
        z_dim=2,
        dim_mult=[1, 2],
        num_res_blocks=1,
        temperal_downsample=[True],
        scale_factor_spatial=2,
        scale_factor_temporal=2,
    ).eval()


@pytest.mark.parametrize("patch_size", [None, 2])
@pytest.mark.parametrize("return_dict", [False, True])
@torch.no_grad()
def test_chunked_decode_matches_diffusers(vae, patch_size, return_dict):
    if patch_size is not None:
        vae = OmniAutoencoderKLWan.from_config(vae.config, patch_size=patch_size, out_channels=12).eval()
    z = torch.randn(1, 2, 3, 2, 2)
    expected = AutoencoderKLWan._decode(vae, z).sample
    actual = vae.decode(z, return_dict=return_dict)
    torch.testing.assert_close(actual.sample if return_dict else actual[0], expected, rtol=0, atol=0)
    chunks: list[torch.Tensor] = []
    assert supports_chunked_vae_decode(vae)
    assert not supports_chunked_vae_decode(object())
    assert vae.decode_with_chunks(z, on_chunk=chunks.append) is None
    assert [chunk.shape[2] for chunk in chunks] == [1, 2, 2]
    torch.testing.assert_close(torch.cat(chunks, dim=2), expected, rtol=0, atol=0)
    assert all(entry is None for entry in vae._feat_map)


@torch.no_grad()
def test_slicing_and_forward_hooks(vae):
    z = torch.randn(2, 2, 2, 2, 2)
    vae.enable_slicing()
    batches = []
    handle = vae.post_quant_conv.register_forward_pre_hook(lambda module, args: batches.append(args[0].shape[0]))
    vae._hf_hook = SimpleNamespace(pre_forward=Mock())
    try:
        assert vae.decode(z).sample.shape[0] == 2
        assert batches == [1, 1]
        vae._hf_hook.pre_forward.assert_called_once_with(vae)
        batches.clear()
        with pytest.raises(ValueError, match="use_slicing"):
            vae.decode_with_chunks(z, on_chunk=Mock())
        assert batches == []
        vae._hf_hook.pre_forward.reset_mock()
        vae.decode_with_chunks(z[:1], on_chunk=Mock())
        assert batches == [1]
        vae._hf_hook.pre_forward.assert_called_once_with(vae)
    finally:
        handle.remove()


@torch.no_grad()
def test_callback_failure_drains_and_clears_cache(vae):
    z = torch.randn(1, 2, 4, 2, 2)
    decoded = Mock(return_value=None)
    handle = vae.decoder.register_forward_hook(lambda *args: decoded())
    error = RuntimeError("consumer failed")
    callback = Mock(side_effect=error)
    try:
        with pytest.raises(RuntimeError, match="consumer failed") as exc:
            vae.decode_with_chunks(z, on_chunk=callback)
        assert exc.value is error
        assert callback.call_count == 1
        assert decoded.call_count == z.shape[2]
        assert all(entry is None for entry in vae._feat_map)
    finally:
        handle.remove()
    # The same instance can decode again after the failure.
    chunks: list[torch.Tensor] = []
    vae.decode_with_chunks(z, on_chunk=chunks.append)
    torch.testing.assert_close(torch.cat(chunks, 2), vae.decode(z).sample, rtol=0, atol=0)


def test_chunked_decode_rejects_spatial_tiling(vae):
    vae.enable_tiling(tile_sample_min_height=2, tile_sample_min_width=2)
    with pytest.raises(ValueError, match="spatial tiling"):
        vae.decode_with_chunks(torch.zeros(1, 2, 2, 2, 2), on_chunk=Mock())


@pytest.mark.parametrize("rank", [0, 1])
@pytest.mark.parametrize("fail_callback", [False, True])
def test_spatial_callback_ownership_and_bounded_retention(monkeypatch, rank, fail_callback):
    # Isolate the temporal loop from halo kernels; weakrefs detect retention of
    # decoded GPU-sized tensors even after the consumer has failed.
    monkeypatch.setattr(wan_spatial_shard, "install_wan_spatial_shard_decode", lambda *args, **kwargs: None)
    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (rank, 2))
    references: list[weakref.ReferenceType[torch.Tensor]] = []

    def decode(x, **kwargs):
        assert sum(ref() is not None for ref in references) <= 1
        chunk = x.clone()
        references.append(weakref.ref(chunk))
        return chunk

    vae = SimpleNamespace(
        config=SimpleNamespace(patch_size=None),
        clear_cache=Mock(),
        post_quant_conv=lambda z: z,
        decoder=decode,
        _feat_map=[],
    )
    callback = Mock(side_effect=RuntimeError("consumer failed") if fail_callback else None)
    z = torch.zeros(1, 3, 6, 2, 2)
    if rank == 0 and fail_callback:
        with pytest.raises(RuntimeError, match="consumer failed"):
            wan_spatial_shard.spatial_shard_decode(vae, z, group=None, on_chunk=callback)
    else:
        assert wan_spatial_shard.spatial_shard_decode(vae, z, group=None, on_chunk=callback) is None
    assert len(references) == z.shape[2]
    assert callback.call_count == (0 if rank else 1 if fail_callback else z.shape[2])
    assert vae.clear_cache.call_count == 2


@pytest.mark.parametrize("rank", [0, 1])
@torch.no_grad()
def test_distributed_small_latent_callback_ownership(vae, monkeypatch, rank):
    vae.is_distributed_enabled = lambda: True
    vae.distributed_executor = SimpleNamespace(group=None)
    monkeypatch.setattr(torch.distributed, "get_rank", lambda **kwargs: rank)
    callback = Mock()
    vae.decode_with_chunks(torch.zeros(1, 2, 2, 2, 2), on_chunk=callback)
    assert callback.call_count == (2 if rank == 0 else 0)


@pytest.mark.parametrize("mode", ["spatial_shard_height", "spatial_shard_width", "tile"])
def test_distributed_chunked_routing(monkeypatch, mode):
    vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan)
    torch.nn.Module.__init__(vae)
    vae.use_tiling = True
    vae.use_slicing = False
    vae.tile_sample_min_height = vae.tile_sample_min_width = 2
    vae.spatial_compression_ratio = 2
    vae.distributed_executor = SimpleNamespace(group=None, parallel_mode=mode, parallel_size=2)
    monkeypatch.setattr(vae, "is_distributed_enabled", lambda: True)
    monkeypatch.setattr(torch.distributed, "get_world_size", lambda **kwargs: 2)
    decode = Mock()
    monkeypatch.setattr(wan_spatial_shard, "spatial_shard_decode", decode)
    callback = Mock()
    z = torch.zeros(1, 2, 2, 2, 2)
    if mode == "tile":
        with pytest.raises(ValueError, match="require spatial-shard"):
            vae.decode_with_chunks(z, on_chunk=callback)
        decode.assert_not_called()
    else:
        assert vae.decode_with_chunks(z, on_chunk=callback) is None
        decode.assert_called_once_with(
            vae, z, group=None, return_dict=True, split_dim=mode.removeprefix("spatial_shard_"), on_chunk=callback
        )