# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from unittest.mock import MagicMock, patch

import pytest
import torch

from vllm_omni.diffusion.models.flux.pipeline_flux import FluxDMD2Pipeline, FluxPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
    LTX2I2VDMD2Pipeline,
    LTX2Pipeline,
    LTX2T2VDMD2Pipeline,
)
from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import QwenImageDMD2Pipeline, QwenImagePipeline
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline, WanT2VDMD2Pipeline
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import Wan22I2VPipeline, WanI2VDMD2Pipeline
from vllm_omni.diffusion.request import OmniDiffusionRequest, OmniDiffusionSamplingParams
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch

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

# DMD2 subclass → immediate base pipeline whose __init__ loads model weights (mocked in tests).
_DMD2_BASE = {
    WanT2VDMD2Pipeline: Wan22Pipeline,
    WanI2VDMD2Pipeline: Wan22I2VPipeline,
    LTX2T2VDMD2Pipeline: LTX2Pipeline,
    LTX2I2VDMD2Pipeline: LTX2Pipeline,
    FluxDMD2Pipeline: FluxPipeline,
    QwenImageDMD2Pipeline: QwenImagePipeline,
}


def _make_pipeline(cls):
    """Run the DMD2 __init__ with the base pipeline mocked out (no model weights loaded)."""

    base = _DMD2_BASE[cls]
    od_config = MagicMock()
    od_config.model = "/nonexistent"

    def _mock_base_init(self, *a, **kw):
        self.od_config = od_config

    with patch.object(base, "__init__", _mock_base_init):
        pipeline = object.__new__(cls)
        torch.nn.Module.__init__(pipeline)
        cls.__init__(pipeline, od_config=od_config)
    return pipeline


def _make_request(prompt=None, **sp_kwargs) -> OmniDiffusionRequest:
    sp = OmniDiffusionSamplingParams(**sp_kwargs)
    return OmniDiffusionRequest(
        prompt=prompt if prompt is not None else {"prompt": "a cat dancing"},
        sampling_params=sp,
        request_id="dmd2-sanitize",
    )


@pytest.fixture(
    params=list(_DMD2_BASE.keys()),
    ids=["wan_t2v", "wan_i2v", "ltx2_t2v", "ltx2_i2v", "flux", "qwen_image"],
)
def pipeline(request):
    return _make_pipeline(request.param)


# ---------------------------------------------------------------------------
# num_inference_steps
# ---------------------------------------------------------------------------


def test_num_inference_steps_forced_to_dmd2_value(pipeline):
    req = _make_request(num_inference_steps=40)
    pipeline._sanitize_dmd2_request(req)
    assert req.sampling_params.num_inference_steps == pipeline.dmd2_config.num_inference_steps


def test_num_inference_steps_already_correct(pipeline):
    req = _make_request(num_inference_steps=pipeline.dmd2_config.num_inference_steps)
    pipeline._sanitize_dmd2_request(req)
    assert req.sampling_params.num_inference_steps == pipeline.dmd2_config.num_inference_steps


# ---------------------------------------------------------------------------
# guidance_scale
# ---------------------------------------------------------------------------


def test_guidance_scale_forced_to_one(pipeline):
    req = _make_request(guidance_scale=5.0, guidance_scale_provided=True)
    pipeline._sanitize_dmd2_request(req)
    assert req.sampling_params.guidance_scale == pipeline.dmd2_config.guidance_scale
    assert req.sampling_params.guidance_scale_provided is False


def test_guidance_scale_already_correct(pipeline):
    req = _make_request(guidance_scale=pipeline.dmd2_config.guidance_scale, guidance_scale_provided=False)
    pipeline._sanitize_dmd2_request(req)
    assert req.sampling_params.guidance_scale == pipeline.dmd2_config.guidance_scale


def test_guidance_scale_provided_flag_cleared(pipeline):
    """guidance_scale_provided=True must be cleared even if scale is already dmd2_guidance_scale."""
    req = _make_request(guidance_scale=pipeline.dmd2_config.guidance_scale, guidance_scale_provided=True)
    pipeline._sanitize_dmd2_request(req)
    assert req.sampling_params.guidance_scale_provided is False


def test_guidance_scale_2_cleared(pipeline):
    req = _make_request(guidance_scale_2=3.0)
    pipeline._sanitize_dmd2_request(req)
    assert req.sampling_params.guidance_scale_2 is None


def test_guidance_scale_2_unset_unchanged(pipeline):
    req = _make_request()
    pipeline._sanitize_dmd2_request(req)
    assert req.sampling_params.guidance_scale_2 is None


def test_true_cfg_scale_cleared(pipeline):
    req = _make_request(true_cfg_scale=2.0)
    pipeline._sanitize_dmd2_request(req)
    assert req.sampling_params.true_cfg_scale is None


def test_do_classifier_free_guidance_forced_false(pipeline):
    req = _make_request(do_classifier_free_guidance=True)
    pipeline._sanitize_dmd2_request(req)
    assert req.sampling_params.do_classifier_free_guidance is False


def test_is_cfg_negative_forced_false(pipeline):
    req = _make_request(is_cfg_negative=True)
    pipeline._sanitize_dmd2_request(req)
    assert req.sampling_params.is_cfg_negative is False


def test_negative_prompt_stripped_from_prompt_dict(pipeline):
    req = _make_request(prompt={"prompt": "a cat", "negative_prompt": "blurry"})
    pipeline._sanitize_dmd2_request(req)
    assert "negative_prompt" not in req.prompt
    assert req.prompt["prompt"] == "a cat"


def test_no_negative_prompt_unchanged(pipeline):
    req = _make_request(prompt={"prompt": "a cat"})
    pipeline._sanitize_dmd2_request(req)
    assert req.prompt == {"prompt": "a cat"}


def test_string_prompt_not_mutated(pipeline):
    """String prompts (not dicts) must pass through unchanged."""
    req = _make_request(prompt="a cat dancing")
    pipeline._sanitize_dmd2_request(req)
    assert req.prompt == "a cat dancing"


def test_request_batch_sanitizes_each_request(pipeline):
    requests = [
        _make_request(
            prompt={"prompt": f"prompt-{index}", "negative_prompt": "blurry"},
            num_inference_steps=40 + index,
            guidance_scale=5.0 + index,
            guidance_scale_provided=True,
        )
        for index in range(2)
    ]

    pipeline._sanitize_dmd2_request(DiffusionRequestBatch(requests))

    for request in requests:
        assert request.sampling_params.num_inference_steps == pipeline.dmd2_config.num_inference_steps
        assert request.sampling_params.guidance_scale == pipeline.dmd2_config.guidance_scale
        assert request.sampling_params.guidance_scale_provided is False
        assert request.sampling_params.do_classifier_free_guidance is False
        assert request.sampling_params.is_cfg_negative is False
        assert "negative_prompt" not in request.prompt


# ---------------------------------------------------------------------------
# Clean request — nothing changes
# ---------------------------------------------------------------------------


def test_sample_solver_stripped_from_extra_args(pipeline):
    """[C1] defense: sample_solver must not leak into req for the base pipeline to read."""
    req = _make_request()
    req.sampling_params.extra_args = {"sample_solver": "euler"}
    pipeline._sanitize_dmd2_request(req)
    assert "sample_solver" not in req.sampling_params.extra_args


def test_flow_shift_stripped_from_extra_args(pipeline):
    """[C1] defense: flow_shift must not leak into req for the base pipeline to read."""
    req = _make_request()
    req.sampling_params.extra_args = {"flow_shift": 3.0}
    pipeline._sanitize_dmd2_request(req)
    assert "flow_shift" not in req.sampling_params.extra_args


def test_unrelated_extra_args_preserved(pipeline):
    """Sanitizer only strips sample_solver / flow_shift; other extras pass through."""
    req = _make_request()
    req.sampling_params.extra_args = {"sample_solver": "euler", "unrelated": 42}
    pipeline._sanitize_dmd2_request(req)
    assert req.sampling_params.extra_args == {"unrelated": 42}


# ---------------------------------------------------------------------------
# Clean request — nothing changes
# ---------------------------------------------------------------------------


def test_clean_request_no_changes(pipeline):
    req = _make_request(
        guidance_scale=pipeline.dmd2_config.guidance_scale,
        guidance_scale_provided=False,
        do_classifier_free_guidance=False,
        is_cfg_negative=False,
    )
    pipeline._sanitize_dmd2_request(req)
    assert req.sampling_params.guidance_scale == pipeline.dmd2_config.guidance_scale
    assert req.sampling_params.guidance_scale_provided is False
    assert req.sampling_params.guidance_scale_2 is None
    assert req.sampling_params.true_cfg_scale is None
    assert req.sampling_params.do_classifier_free_guidance is False
    assert req.sampling_params.is_cfg_negative is False