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

from types import SimpleNamespace
from unittest.mock import Mock

import pytest

from vllm_omni.diffusion import io_support
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.inputs.data import OmniDiffusionSamplingParams

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


@pytest.mark.parametrize("model_class_name", ["DreamZeroPipeline", "LingBotWorldCausalDMDPipeline"])
def test_observation_conditioned_startup_skips_generic_warmup(model_class_name: str) -> None:
    """Observation-conditioned pipelines must not receive generic text requests."""
    engine = DiffusionEngine.__new__(DiffusionEngine)
    engine.od_config = OmniDiffusionConfig.__new__(OmniDiffusionConfig)
    engine.od_config.model_class_name = model_class_name
    engine.od_config.diffusion_load_format = "default"
    engine.add_req_and_wait_for_response = Mock(side_effect=AssertionError("generic text warmup submitted"))
    engine.close = Mock()

    engine.run_startup_warmup()

    engine.add_req_and_wait_for_response.assert_not_called()
    engine.close.assert_not_called()


@pytest.mark.parametrize("step_execution", [False, True])
def test_dummy_run_uses_enough_steps_for_execution_mode(
    monkeypatch: pytest.MonkeyPatch,
    step_execution: bool,
) -> None:
    engine = DiffusionEngine.__new__(DiffusionEngine)
    engine.od_config = type(
        "Config",
        (),
        {"model_class_name": "mock_model", "diffusion_load_format": "default"},
    )()
    engine.step_execution = step_execution
    engine.pre_process_func = None
    captured_requests = []

    monkeypatch.setattr(
        "vllm_omni.diffusion.diffusion_engine.supports_multimodal_input",
        lambda od_config: (False, False),
    )
    monkeypatch.setattr(
        "vllm_omni.diffusion.diffusion_engine.get_dummy_run_num_frames",
        lambda model_class_name, supports_audio_input: 1,
    )

    def run_request(request):
        captured_requests.append(request)
        return type("Output", (), {"error": None})()

    engine.add_req_and_wait_for_response = run_request

    engine._dummy_run()

    assert len(captured_requests) == 1
    assert captured_requests[0].sampling_params.num_inference_steps == 2


def test_allgather_startup_runs_broadcast_dummy_request() -> None:
    engine = object.__new__(DiffusionEngine)
    engine.od_config = SimpleNamespace(
        diffusion_offload_config={
            "mode": "layer",
            "components": ["dit"],
            "layer_options": {"dit": {"weight_transfer": "allgather"}},
        },
        parallel_config=SimpleNamespace(data_parallel_size=2, sequence_parallel_size=1),
    )
    engine._dummy_run = Mock()

    engine.run_startup_warmup()

    engine._dummy_run.assert_called_once_with()


def test_dummy_run_num_frames_uses_explicit_model_setting(monkeypatch: pytest.MonkeyPatch) -> None:
    class JointAudioVideoModel:
        dummy_run_num_frames = 2

    monkeypatch.setattr(
        io_support.DiffusionModelRegistry,
        "_try_load_model_cls",
        lambda model_class_name: JointAudioVideoModel,
    )

    assert io_support.get_dummy_run_num_frames("joint_audio_video", supports_audio_input=False) == 2


def test_dummy_run_num_frames_keeps_audio_output_default(monkeypatch: pytest.MonkeyPatch) -> None:
    class AudioOutputModel:
        support_audio_output = True

    monkeypatch.setattr(
        io_support.DiffusionModelRegistry,
        "_try_load_model_cls",
        lambda model_class_name: AudioOutputModel,
    )

    assert io_support.get_dummy_run_num_frames("audio_output", supports_audio_input=False) == 2


def test_dummy_run_num_frames_defaults_to_single_frame(monkeypatch: pytest.MonkeyPatch) -> None:
    class VideoOnlyModel:
        pass

    monkeypatch.setattr(
        io_support.DiffusionModelRegistry,
        "_try_load_model_cls",
        lambda model_class_name: VideoOnlyModel,
    )

    assert io_support.get_dummy_run_num_frames("video_only", supports_audio_input=False) == 1


def test_dummy_run_num_frames_uses_audio_input_fallback(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(
        io_support.DiffusionModelRegistry,
        "_try_load_model_cls",
        lambda model_class_name: None,
    )

    assert io_support.get_dummy_run_num_frames("unknown", supports_audio_input=True) == 2


def test_dummy_run_image_count_resolves_hunyuan_architecture_alias() -> None:
    assert io_support.get_dummy_run_num_image_inputs("HunyuanImage3ForCausalMM") == 3
    assert io_support.get_dummy_run_num_image_inputs("unknown") == 1


def test_dense_mode_does_not_build_kv_profile_request() -> None:
    engine = object.__new__(DiffusionEngine)
    engine.od_config = SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY)
    engine._make_dummy_request = Mock(side_effect=AssertionError("dense mode must not profile"))

    assert engine._prepare_diffusion_kv_profile_requests() is None
    engine._make_dummy_request.assert_not_called()


@pytest.mark.parametrize(
    ("execution_mode", "uses_dlo_dp", "max_num_seqs", "expected_profile_requests"),
    [
        (DiffusionExecutionMode.STEP_BATCH, False, 2, 2),
        (DiffusionExecutionMode.REQUEST_BATCH, True, 4, 1),
        (DiffusionExecutionMode.STEP_BATCH, True, 3, 3),
    ],
)
def test_paged_kv_profile_requests_match_per_rank_batch(
    execution_mode: DiffusionExecutionMode,
    uses_dlo_dp: bool,
    max_num_seqs: int,
    expected_profile_requests: int,
) -> None:
    engine = object.__new__(DiffusionEngine)
    engine.execution_mode = execution_mode
    engine.od_config = SimpleNamespace(
        diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
        model_class_name="HunyuanImage3ForCausalMM",
        max_num_seqs=max_num_seqs,
        parallel_config=SimpleNamespace(data_parallel_size=max_num_seqs if uses_dlo_dp else 1),
        enable_distributed_layerwise_offload=uses_dlo_dp,
        dlo_use_allgather=True,
    )
    request = OmniDiffusionRequest(
        prompt="profile",
        request_id="profile-request",
        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
    )
    engine._make_dummy_request = Mock(return_value=request)
    prepared_layout = object()
    scheduler_kv_state = (
        DiffusionKVRequest(
            "profile-request/diffusion-kv/0",
            sequence_id=0,
            prefix_len=3,
            target_len=8,
            seq_len=12,
        ),
        DiffusionKVRequest(
            "profile-request/diffusion-kv/1",
            sequence_id=1,
            prefix_len=4,
            target_len=8,
            seq_len=13,
        ),
    )

    def preprocess(req: OmniDiffusionRequest) -> OmniDiffusionRequest:
        req.prepared_layout = prepared_layout
        req.diffusion_kv_requests = scheduler_kv_state
        return req

    engine._prepare_request_for_admission = Mock(side_effect=preprocess)

    result = engine._prepare_diffusion_kv_profile_requests()

    assert result is not None
    assert len(result) == expected_profile_requests
    assert len({profile_request.request_id for profile_request in result}) == expected_profile_requests
    assert all(profile_request.is_dummy_run() for profile_request in result)
    assert all(profile_request.prepared_layout is prepared_layout for profile_request in result)
    assert all(profile_request.diffusion_kv_requests is None for profile_request in result)
    assert len({id(profile_request.sampling_params) for profile_request in result}) == expected_profile_requests
    assert engine._diffusion_kv_profile_limits == (2, 13, 8)
    assert len(result) * engine._diffusion_kv_profile_limits[0] == expected_profile_requests * 2
    engine._make_dummy_request.assert_called_once_with(
        height=1024,
        width=1024,
        guidance_scale=5.0,
        num_image_inputs=3,
    )
    engine._prepare_request_for_admission.assert_called_once_with(request)


@pytest.mark.parametrize(
    ("num_sequences", "seq_len", "target_len"),
    [
        (3, 16, 8),
        (1, 17, 8),
        (1, 16, 9),
    ],
)
def test_paged_kv_admission_rejects_shape_beyond_profile_envelope(
    num_sequences: int,
    seq_len: int,
    target_len: int,
) -> None:
    engine = object.__new__(DiffusionEngine)
    engine._diffusion_kv_profile_limits = (2, 16, 8)
    engine.pre_process_func = lambda request: request
    request = OmniDiffusionRequest(
        prompt="too large",
        request_id="too-large",
        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
        diffusion_kv_requests=tuple(
            DiffusionKVRequest(
                f"too-large/diffusion-kv/{sequence_id}",
                sequence_id=sequence_id,
                prefix_len=0,
                target_len=target_len,
                seq_len=seq_len,
            )
            for sequence_id in range(num_sequences)
        ),
    )

    with pytest.raises(ValueError, match="exceeds the startup memory-profile envelope"):
        engine._prepare_request_for_admission(request)


def test_paged_kv_admission_accepts_shape_at_profile_envelope() -> None:
    engine = object.__new__(DiffusionEngine)
    engine._diffusion_kv_profile_limits = (2, 16, 8)
    engine.pre_process_func = lambda request: request
    request = OmniDiffusionRequest(
        prompt="fits",
        request_id="fits",
        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
        diffusion_kv_requests=tuple(
            DiffusionKVRequest(
                f"fits/diffusion-kv/{sequence_id}",
                sequence_id=sequence_id,
                prefix_len=8,
                target_len=8,
                seq_len=16,
            )
            for sequence_id in range(2)
        ),
    )

    assert engine._prepare_request_for_admission(request) is request