# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""WAN reference coverage for the typed pre-D2H media contract."""

from __future__ import annotations

import numpy as np
import pytest
import torch

from tests.helpers.mark import hardware_marks
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig, VideoOutputTransportConfig
from vllm_omni.diffusion.ipc import pack_diffusion_output_shm, unpack_diffusion_output_shm
from vllm_omni.diffusion.media import (
    DiffusionMediaOutput,
    VideoMediaOutput,
    VideoTensorEncoding,
    VideoTensorLayout,
    VideoTensorSpec,
    VideoValueRange,
)
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import get_wan22_post_process_func
from vllm_omni.diffusion.postprocess.device_reduction import prepare_diffusion_media_for_transport
from vllm_omni.diffusion.postprocess.media import finalize_diffusion_media
from vllm_omni.entrypoints.openai.video_api_utils import _coerce_video_to_uint8_frames
from vllm_omni.inputs.data import OmniDiffusionSamplingParams

pytestmark = [
    pytest.mark.core_model,
    pytest.mark.diffusion,
    *hardware_marks(res={"cuda": "L4"}, num_cards=1),
]

_GPU = torch.accelerator.is_available() if hasattr(torch, "accelerator") else torch.cuda.is_available()
requires_gpu = pytest.mark.skipif(not _GPU, reason="device reduction is a GPU path")


def _decoded_wan_video(tensor: torch.Tensor) -> DiffusionMediaOutput:
    return DiffusionMediaOutput(
        video=VideoMediaOutput(
            tensor=tensor,
            spec=VideoTensorSpec(
                layout=VideoTensorLayout.BCTHW,
                encoding=VideoTensorEncoding.NORMALIZED_FLOAT,
                value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE,
            ),
        )
    )


def _config(*, enabled: bool) -> OmniDiffusionConfig:
    return OmniDiffusionConfig(
        model=None,
        video_output_transport=VideoOutputTransportConfig(enable_device_postprocess=enabled),
    )


def _encoder_frames(video_payload: np.ndarray) -> list[np.ndarray]:
    return [_coerce_video_to_uint8_frames(video_payload[i]) for i in range(video_payload.shape[0])]


@requires_gpu
def test_wan_media_reduction_matches_float32_reference_and_bounds_native_path() -> None:
    torch.manual_seed(0)
    vae_out = torch.rand(2, 3, 8, 64, 96, device="cuda:0", dtype=torch.bfloat16) * 2.4 - 1.2
    sampling = OmniDiffusionSamplingParams(output_type=None, enable_frame_interpolation=False)

    legacy_postprocess = get_wan22_post_process_func(_config(enabled=False))
    widened_out = legacy_postprocess(vae_out.float(), output_type="np", sampling_params=sampling)
    widened_frames = _encoder_frames(widened_out["payload"]["video"])
    native_out = legacy_postprocess(vae_out, output_type="np", sampling_params=sampling)
    native_frames = _encoder_frames(native_out["payload"]["video"])

    prepared = prepare_diffusion_media_for_transport(
        _decoded_wan_video(vae_out),
        od_config=_config(enabled=True),
        sampling_params=sampling,
    )
    reduced_out = finalize_diffusion_media(prepared, sampling_params=sampling)
    assert prepared.video.spec.encoding is VideoTensorEncoding.UINT8_FRAMES
    assert reduced_out["payload"]["video"].dtype == np.uint8
    reduced_frames = _encoder_frames(reduced_out["payload"]["video"])

    assert len(widened_frames) == len(reduced_frames) == 2
    for expected, native, produced in zip(widened_frames, native_frames, reduced_frames, strict=True):
        assert produced.shape == expected.shape == (8, 64, 96, 3)
        np.testing.assert_array_equal(produced, expected)
        deviation = np.abs(produced.astype(np.int16) - native.astype(np.int16))
        assert deviation.max() <= 1


@requires_gpu
def test_wan_uint8_media_finalization_does_not_rescale() -> None:
    frames = torch.randint(0, 256, (1, 4, 16, 16, 3), dtype=torch.uint8, device="cuda:0")
    media = DiffusionMediaOutput(
        video=VideoMediaOutput(
            tensor=frames,
            spec=VideoTensorSpec(
                layout=VideoTensorLayout.BTHWC,
                encoding=VideoTensorEncoding.UINT8_FRAMES,
                value_range=VideoValueRange.ZERO_TO_255,
            ),
        ),
        prepared_for_transport=True,
    )

    out = finalize_diffusion_media(media, sampling_params=OmniDiffusionSamplingParams(output_type="np"))

    np.testing.assert_array_equal(out["payload"]["video"], frames.cpu().numpy())


@requires_gpu
def test_diffusion_output_to_cpu_moves_wan_media_off_device() -> None:
    prepared = prepare_diffusion_media_for_transport(
        _decoded_wan_video(torch.randn(1, 3, 2, 4, 5, device="cuda:0")),
        od_config=_config(enabled=False),
        sampling_params=OmniDiffusionSamplingParams(output_type="np"),
    )
    output = DiffusionOutput(media=prepared, to_cpu=True)

    assert output.media is not None
    assert output.media.video.tensor.device.type == "cpu"


@requires_gpu
def test_small_typed_media_is_on_cpu_before_ipc() -> None:
    sampling = OmniDiffusionSamplingParams(output_type="np")
    prepared = prepare_diffusion_media_for_transport(
        _decoded_wan_video(torch.randn(1, 3, 2, 4, 5, device="cuda:0")),
        od_config=_config(enabled=True),
        sampling_params=sampling,
    )
    output = DiffusionOutput(media=prepared)

    pack_diffusion_output_shm(output)

    assert isinstance(output.media, dict)
    packed_tensor = output.media["video"]["tensor"]
    assert isinstance(packed_tensor, torch.Tensor)
    assert packed_tensor.device.type == "cpu"
    unpack_diffusion_output_shm(output)


@requires_gpu
def test_wan_float_media_finalization_matches_the_legacy_path() -> None:
    video = torch.rand(1, 3, 4, 16, 16, device="cuda:0", dtype=torch.bfloat16) * 2 - 1
    sampling = OmniDiffusionSamplingParams(output_type="np")
    prepared = prepare_diffusion_media_for_transport(
        _decoded_wan_video(video),
        od_config=_config(enabled=False),
        sampling_params=sampling,
    )

    produced = finalize_diffusion_media(prepared, sampling_params=sampling)
    expected = get_wan22_post_process_func(_config(enabled=False))(
        video,
        output_type="np",
        sampling_params=sampling,
    )

    assert prepared.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT
    np.testing.assert_array_equal(produced["payload"]["video"], expected["payload"]["video"])


@requires_gpu
def test_oom_fallback_media_packs_without_device_widening(monkeypatch: pytest.MonkeyPatch) -> None:
    # Regression: when the fp32 reduction OOMs, the fallback keeps the bf16 float
    # payload. Async IPC must copy it to the host in bf16 and widen on the CPU;
    # a second device-side fp32 conversion would OOM again and drop the payload.
    from vllm_omni.diffusion.postprocess import device_reduction

    # Large enough (> 1 MB in bf16) that packing routes the tensor through the
    # shared-memory D2H path rather than the small-tensor CPU copy.
    video = torch.rand(1, 3, 16, 128, 128, device="cuda:0", dtype=torch.bfloat16) * 2 - 1

    def _raise_oom(*_args: object, **_kwargs: object) -> torch.Tensor:
        raise torch.OutOfMemoryError("forced fp32 reduction OOM")

    monkeypatch.setattr(device_reduction, "reduce_video_to_uint8_frames", _raise_oom)

    prepared = prepare_diffusion_media_for_transport(
        _decoded_wan_video(video),
        od_config=_config(enabled=True),
        sampling_params=OmniDiffusionSamplingParams(output_type="np"),
    )
    assert prepared.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT
    assert prepared.video.tensor.dtype == torch.bfloat16

    output = DiffusionOutput(media=prepared)

    pinned_dtypes: list[torch.dtype] = []
    real_empty = torch.empty

    def _spy_empty(*args: object, **kwargs: object) -> torch.Tensor:
        tensor = real_empty(*args, **kwargs)
        if kwargs.get("pin_memory"):
            pinned_dtypes.append(tensor.dtype)
        return tensor

    monkeypatch.setattr(torch, "empty", _spy_empty)

    d2h_stream = torch.Stream(device=torch.device("cuda", 0))
    pack_diffusion_output_shm(output, d2h_stream=d2h_stream)
    d2h_stream.synchronize()

    assert pinned_dtypes == [torch.bfloat16], (
        "async packing must D2H in the original dtype and widen on the CPU, "
        "never widen bf16 to fp32 on the accelerator before the copy"
    )

    unpack_diffusion_output_shm(output)
    assert isinstance(output.media, DiffusionMediaOutput)
    assert output.media.video.tensor.dtype == torch.bfloat16
    torch.testing.assert_close(output.media.video.tensor, video.cpu())


@requires_gpu
def test_oom_fallback_recovers_the_allocator(monkeypatch: pytest.MonkeyPatch) -> None:
    # A real float32 OOM leaves the CUDA caching allocator in a failed state;
    # the fallback must empty the cache and drain the device before marking the
    # float representation transport-ready, so the worker D2H of the same tensor
    # can actually run. Use a genuine allocation failure (per-process cap) rather
    # than injecting the exception, which never poisons the allocator.
    from vllm_omni.platforms import current_omni_platform

    video = torch.empty(1, 3, 16, 128, 128, device="cuda:0", dtype=torch.bfloat16)
    video.uniform_(-1.0, 1.0)
    expected = video.cpu()
    d2h_stream = torch.Stream(device=video.device)
    base_bytes = video.numel() * video.element_size()
    _free, total = current_omni_platform.get_device_memory()

    # Cap the process so the fp32 conversion (2x base bytes) cannot fit, while
    # retaining the original bf16 tensor. This forces a real allocator OOM at
    # the .to(fp32), rather than injecting an exception from a monkeypatch.
    torch.cuda.set_per_process_memory_fraction(min(0.99, (1.2 * base_bytes) / total))
    real_empty_cache = current_omni_platform.empty_cache
    real_synchronize = current_omni_platform.synchronize
    try:
        empty_calls: list[int] = []
        sync_calls: list[bool] = []

        def _empty_cache() -> None:
            real_empty_cache()
            empty_calls.append(torch.accelerator.current_device_index())

        def _synchronize() -> None:
            real_synchronize()
            sync_calls.append(True)

        monkeypatch.setattr(current_omni_platform, "empty_cache", _empty_cache)
        monkeypatch.setattr(current_omni_platform, "synchronize", _synchronize)

        prepared = prepare_diffusion_media_for_transport(
            _decoded_wan_video(video),
            od_config=_config(enabled=True),
            sampling_params=OmniDiffusionSamplingParams(output_type="np"),
        )

        # The fallback ran past a genuine OOM and kept the bf16 float payload.
        assert prepared.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT
        assert prepared.video.tensor.dtype == torch.bfloat16
        assert empty_calls, "fallback must empty the CUDA allocator cache after a real OOM"
        assert sync_calls, "fallback must drain the device after a real OOM"

        # Exercise the next production step: async D2H/SHM packing of the same
        # tensor must succeed after allocator recovery and preserve its values.
        output = DiffusionOutput(media=prepared)
        pack_diffusion_output_shm(output, d2h_stream=d2h_stream)
        d2h_stream.synchronize()
        unpack_diffusion_output_shm(output)
        assert isinstance(output.media, DiffusionMediaOutput)
        torch.testing.assert_close(output.media.video.tensor, expected)
    finally:
        torch.cuda.set_per_process_memory_fraction(1.0)
        real_empty_cache()
        real_synchronize()


@requires_gpu
def test_oom_fallback_releases_late_uint8_traceback_before_clone(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """A final uint8 allocation OOM must release float intermediates first."""
    from vllm_omni.diffusion.postprocess import device_reduction
    from vllm_omni.platforms import current_omni_platform

    # A split request is a non-compact view, so float fallback must allocate a
    # request-owned clone after reduction fails.
    backing = torch.empty(2, 3, 16, 128, 128, device="cuda:0", dtype=torch.bfloat16)
    backing.uniform_(-1.0, 1.0)
    video = backing[1:2].detach()
    expected = video.cpu()
    _, total = current_omni_platform.get_device_memory()
    real_to = torch.Tensor.to
    real_empty_cache = current_omni_platform.empty_cache
    real_synchronize = current_omni_platform.synchronize
    real_ensure_request_owned = device_reduction.ensure_request_owned_tensor
    uint8_attempts: list[int] = []
    fallback_allocations: list[int] = []

    def _record_fallback_allocation(tensor: torch.Tensor) -> torch.Tensor:
        fallback_allocations.append(torch.accelerator.memory_allocated(tensor.device))
        return real_ensure_request_owned(tensor)

    def _fail_real_uint8_allocation(
        tensor: torch.Tensor,
        *args: object,
        **kwargs: object,
    ) -> torch.Tensor:
        dtype = kwargs.get("dtype")
        if dtype is None and args and isinstance(args[0], torch.dtype):
            dtype = args[0]
        if tensor.is_cuda and dtype is torch.uint8:
            # Earlier float32 work runs without a cap. Immediately before the
            # real final uint8 allocation, cap the allocator below that request
            # but above the post-unwind request-owned bf16 clone requirement.
            real_empty_cache()
            allocated = torch.accelerator.memory_allocated(tensor.device)
            uint8_bytes = tensor.numel()
            limit_bytes = allocated + max(1, uint8_bytes // 2)
            torch.cuda.set_per_process_memory_fraction(min(0.99, limit_bytes / total))
            uint8_attempts.append(allocated)
        return real_to(tensor, *args, **kwargs)

    monkeypatch.setattr(device_reduction, "ensure_request_owned_tensor", _record_fallback_allocation)
    monkeypatch.setattr(torch.Tensor, "to", _fail_real_uint8_allocation)
    try:
        prepared = prepare_diffusion_media_for_transport(
            _decoded_wan_video(video),
            od_config=_config(enabled=True),
            sampling_params=OmniDiffusionSamplingParams(output_type="np"),
        )

        assert uint8_attempts, "the real final .to(uint8) allocation was not attempted"
        assert fallback_allocations, "float fallback did not request compact storage"
        assert fallback_allocations[0] < uint8_attempts[0], (
            "the failed reduction's float intermediates remained live during fallback"
        )
        assert prepared.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT
        assert prepared.video.tensor.dtype is torch.bfloat16
        assert prepared.video.tensor.storage_offset() == 0
        assert prepared.video.tensor.untyped_storage().data_ptr() != video.untyped_storage().data_ptr()
        torch.testing.assert_close(prepared.video.tensor.cpu(), expected)
    finally:
        torch.cuda.set_per_process_memory_fraction(1.0)
        real_empty_cache()
        real_synchronize()