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

"""
System test for cache-dit backend.

This test verifies that cache-dit acceleration works correctly with diffusion models.
It uses minimal settings to keep test time short for CI.
"""

import pytest
import torch

from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniRunner
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform

# Use random weights model for testing
models = ["riverclouds/qwen_image_random"]


@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cache
@hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"})
@pytest.mark.parametrize("model_name", models)
def test_cache_dit(model_name: str):
    """Test cache-dit backend with diffusion model."""
    # Configure cache-dit with minimal settings for fast testing
    cache_config = {
        "Fn_compute_blocks": 1,
        "Bn_compute_blocks": 0,
        "max_warmup_steps": 2,  # Minimal warmup for fast test
        "residual_diff_threshold": 0.24,
        "max_continuous_cached_steps": 3,
    }
    with OmniRunner(
        model_name,
        cache_backend="cache_dit",
        cache_config=cache_config,
    ) as runner:
        # Use minimal settings for fast testing
        height = 256
        width = 256
        num_inference_steps = 4  # Minimal steps for fast test

        outputs = runner.omni.generate(
            "a photo of a cat sitting on a laptop keyboard",
            OmniDiffusionSamplingParams(
                height=height,
                width=width,
                num_inference_steps=num_inference_steps,
                guidance_scale=0.0,
                generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
                num_outputs_per_prompt=1,  # Single output for speed
            ),
        )
        # Extract images from request_output['images']
        first_output = outputs[0]
        assert first_output.final_output_type == "image"
        if not hasattr(first_output, "request_output") or not first_output.request_output:
            raise ValueError("No request_output found in OmniRequestOutput")

        req_out = first_output.request_output
        if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
            raise ValueError("Invalid request_output structure or missing 'images' key")

        images = req_out.images

        # Verify generation succeeded
        assert images is not None
        assert len(images) == 1
        # Check image size
        assert images[0].width == width
        assert images[0].height == height