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

"""
End-to-end test for SenseNova-U1 img2img generation.

This test validates that the SenseNova-U1 model generates edited images from
an input image and text prompt that match expected reference pixel values
within a ±10 tolerance.

Equivalent to running:
    python examples/offline_inference/sensenova_u1/end2end.py \
        --prompt "Turn this into an oil painting" \
        --image 2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg \
        --width 2560 --height 1669 \
        --seed 42 --num-steps 50 \
        --cfg-scale 4.0 --img-cfg-scale 1.0 --cfg-norm none \
        --think --print-think
"""

import os
from typing import Any

os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"

import pytest
from PIL import Image
from vllm.assets.image import ImageAsset

from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniRunner
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams

# Reference pixel data extracted from the known-good output image generated by:
#   python examples/offline_inference/sensenova_u1/end2end.py \
#       --prompt "Turn this into an oil painting" \
#       --image 2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg \
#       --width 2560 --height 1669 --seed 42 --num-steps 50 \
#       --cfg-scale 4.0 --img-cfg-scale 1.0 --cfg-norm none --think
REFERENCE_PIXELS = [
    {"position": (100, 100), "rgb": (95, 131, 174)},
    {"position": (400, 50), "rgb": (142, 153, 174)},
    {"position": (700, 100), "rgb": (127, 149, 174)},
    {"position": (150, 400), "rgb": (99, 141, 183)},
    {"position": (512, 336), "rgb": (184, 177, 165)},
    {"position": (700, 400), "rgb": (106, 139, 168)},
    {"position": (100, 600), "rgb": (170, 172, 155)},
    {"position": (400, 600), "rgb": (159, 171, 169)},
    {"position": (700, 600), "rgb": (139, 172, 187)},
    {"position": (256, 256), "rgb": (136, 158, 180)},
]

PIXEL_TOLERANCE = 10

DEFAULT_PROMPT = "Turn this into an oil painting"

EXPECTED_OUTPUT_SIZE = (2560, 1664)


def _load_input_image() -> Image.Image:
    """Load the test input image via vllm's ImageAsset."""
    return ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB")


def _build_sampling_params() -> OmniDiffusionSamplingParams:
    """Build sampling parameters for SenseNova-U1 img2img generation."""
    return OmniDiffusionSamplingParams(
        height=1669,
        width=2560,
        seed=42,
        num_inference_steps=50,
        extra_args={
            "cfg_scale": 4.0,
            "img_cfg_scale": 1.0,
            "cfg_norm": "none",
            "timestep_shift": 3.0,
            "cfg_interval": (0.0, 1.0),
            "batch_size": 1,
            "think": True,
            "t_eps": 0.02,
        },
    )


def _extract_generated_image(omni_outputs: list) -> Image.Image | None:
    """Extract the generated image from Omni outputs (single-stage DiT)."""
    for req_output in omni_outputs:
        if images := getattr(req_output, "images", None):
            return images[0]
    return None


def _validate_pixels(
    image: Image.Image,
    reference_pixels: list[dict[str, Any]] = REFERENCE_PIXELS,
    tolerance: int = PIXEL_TOLERANCE,
) -> None:
    """Validate that image pixels match expected reference values.

    Args:
        image: The PIL Image to validate.
        reference_pixels: List of dicts with 'position' (x, y) and 'rgb' (R, G, B).
        tolerance: Maximum allowed difference per color channel.

    Raises:
        AssertionError: If any pixel differs beyond tolerance.
    """
    for ref in reference_pixels:
        x, y = ref["position"]
        expected = ref["rgb"]
        actual = image.getpixel((x, y))[:3]
        assert all(abs(a - e) <= tolerance for a, e in zip(actual, expected)), (
            f"Pixel mismatch at ({x}, {y}): expected {expected}, got {actual}"
        )


def _generate_sensenova_u1_img2img(
    omni: Omni,
    input_image: Image.Image,
    prompt: str = DEFAULT_PROMPT,
) -> Image.Image:
    """Generate an edited image using SenseNova-U1 model with img2img pipeline.

    Args:
        omni: The Omni instance to use for generation.
        input_image: The input PIL Image for img2img.
        prompt: The text prompt for image editing.

    Returns:
        The generated PIL Image.

    Raises:
        AssertionError: If no image is generated or size is incorrect.
    """
    sampling_params = _build_sampling_params()

    omni_outputs = list(
        omni.generate(
            prompts={
                "prompt": prompt,
                "multi_modal_data": {"image": input_image},
                "modalities": ["img2img"],
            },
            sampling_params_list=sampling_params,
        )
    )

    generated_image = _extract_generated_image(omni_outputs)
    assert generated_image is not None, "No images generated"
    assert generated_image.size == EXPECTED_OUTPUT_SIZE, f"Expected {EXPECTED_OUTPUT_SIZE}, got {generated_image.size}"

    return generated_image


@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"})
def test_sensenova_u1_img2img(run_level):
    """Test SenseNova-U1 img2img (single-stage diffusion, no deploy YAML)."""
    input_image = _load_input_image()
    with OmniRunner(
        "SenseNova/SenseNova-U1-8B-MoT",
        stage_configs_path=None,
    ) as runner:
        generated_image = _generate_sensenova_u1_img2img(runner.omni, input_image)
        if run_level == "advanced_model":
            _validate_pixels(generated_image)


@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cache
@hardware_test(res={"cuda": "H100"})
def test_sensenova_u1_img2img_cache_dit():
    """Test SenseNova-U1 img2img with Cache-DiT enabled."""
    input_image = _load_input_image()
    with OmniRunner(
        "SenseNova/SenseNova-U1-8B-MoT",
        stage_configs_path=None,
        cache_backend="cache_dit",
    ) as runner:
        _generate_sensenova_u1_img2img(runner.omni, input_image)