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

from types import SimpleNamespace

import pytest

from vllm_omni.diffusion.data import AttentionConfig
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
from vllm_omni.entrypoints.cli.serve import OmniServeCommand
from vllm_omni.utils.tracking_parser import TrackingArgumentParser

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


def test_default_stage_config_includes_cache_backend():
    """Ensure cache knobs survive the default diffusion-stage builder."""
    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
        {
            "cache_backend": "cache_dit",
            "cache_config": '{"Fn_compute_blocks": 2}',
            "vae_use_slicing": True,
            "ulysses_degree": 2,
        }
    )[0]

    engine_args = stage_cfg["engine_args"]
    assert stage_cfg["stage_type"] == "diffusion"
    assert engine_args["cache_backend"] == "cache_dit"
    assert engine_args["cache_config"]["Fn_compute_blocks"] == 2
    assert engine_args["vae_use_slicing"] is True
    assert engine_args["parallel_config"].ulysses_degree == 2
    assert engine_args["model_stage"] == "diffusion"


def test_default_cache_config_used_when_missing():
    """Ensure default cache_config is synthesized when only backend is given."""
    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
        {
            "cache_backend": "cache_dit",
        }
    )[0]

    cache_config = stage_cfg["engine_args"]["cache_config"]
    assert cache_config is not None
    assert cache_config["Fn_compute_blocks"] == 1


def test_default_stage_devices_from_sequence_parallel():
    """Ensure runtime devices reflect computed diffusion world size."""
    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
        {
            "ulysses_degree": 2,
            "ring_degree": 2,
        }
    )[0]

    assert stage_cfg["runtime"]["devices"] == "0,1,2,3"


def test_default_stage_config_uses_parallel_size_kwargs():
    """Ensure default diffusion parallel_config uses CLI/API parallel sizes."""
    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
        {
            "pipeline_parallel_size": 2,
            "data_parallel_size": 3,
            "tensor_parallel_size": 4,
            "enable_expert_parallel": True,
        }
    )[0]

    parallel_config = stage_cfg["engine_args"]["parallel_config"]
    assert parallel_config.pipeline_parallel_size == 2
    assert parallel_config.data_parallel_size == 3
    assert parallel_config.tensor_parallel_size == 4
    assert parallel_config.enable_expert_parallel is True


def test_default_stage_config_defaults_nullified_parallel_size_kwargs():
    """Ensure nullified diffusion parallel-size kwargs fall back to defaults."""
    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
        {
            "pipeline_parallel_size": None,
            "data_parallel_size": None,
            "tensor_parallel_size": None,
            "enable_expert_parallel": None,
            "enforce_eager": None,
        }
    )[0]

    parallel_config = stage_cfg["engine_args"]["parallel_config"]
    assert parallel_config.pipeline_parallel_size == 1
    assert parallel_config.data_parallel_size == 1
    assert parallel_config.tensor_parallel_size == 1
    assert parallel_config.enable_expert_parallel is False
    assert stage_cfg["engine_args"]["enforce_eager"] is False


def test_default_stage_config_propagates_ulysses_mode():
    """Ensure UAA mode survives default diffusion-stage creation."""
    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
        {
            "ulysses_degree": 4,
            "ulysses_mode": "advanced_uaa",
        }
    )[0]

    parallel_config = stage_cfg["engine_args"]["parallel_config"]
    assert parallel_config.ulysses_degree == 4
    assert parallel_config.ulysses_mode == "advanced_uaa"


def test_default_stage_config_includes_default_sampling_params():
    """Ensure default sampling params survive the default diffusion-stage builder."""
    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
        {
            "default_sampling_params": '{"0": {"generator_device":"cpu", "guidance_scale":7.5}}',
        }
    )[0]

    assert stage_cfg["default_sampling_params"] == {
        "generator_device": "cpu",
        "guidance_scale": 7.5,
    }


def test_default_stage_config_includes_diffusion_attention_backend():
    """Ensure diffusion attention shorthand lands in engine_args.diffusion_attention_config."""
    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
        {
            "diffusion_attention_backend": "FLASH_ATTN",
        }
    )[0]

    diffusion_attention_config = stage_cfg["engine_args"]["diffusion_attention_config"]
    assert isinstance(diffusion_attention_config, AttentionConfig)
    assert diffusion_attention_config.default is not None
    assert diffusion_attention_config.default.backend == "FLASH_ATTN"


def test_default_stage_config_includes_diffusion_attention_config():
    """Ensure structured diffusion attention config survives default stage creation."""
    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
        {
            "diffusion_attention_config": {
                "default": {"backend": "FLASH_ATTN"},
                "per_role": {"cross": {"backend": "TORCH_SDPA"}},
            },
        }
    )[0]

    diffusion_attention_config = stage_cfg["engine_args"]["diffusion_attention_config"]
    assert isinstance(diffusion_attention_config, AttentionConfig)
    assert diffusion_attention_config.default is not None
    assert diffusion_attention_config.default.backend == "FLASH_ATTN"
    assert diffusion_attention_config.per_role["cross"].backend == "TORCH_SDPA"


def test_default_stage_config_rejects_conflicting_diffusion_attention_inputs():
    """Ensure shorthand and default.backend stay mutually exclusive."""
    with pytest.raises(ValueError, match="mutually exclusive"):
        AsyncOmniEngine._create_default_diffusion_stage_cfg(
            {
                "diffusion_attention_backend": "FLASH_ATTN",
                "diffusion_attention_config": {
                    "default": {"backend": "TORCH_SDPA"},
                },
            }
        )


def test_default_stage_config_engine_args():
    """Ensure default diffusion-stage builder sets and propagates engine_args."""
    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
        {
            "distributed_executor_backend": "ray",
            "boundary_ratio": 0.875,
            "flow_shift": 5.0,
            "trust_remote_code": True,
        }
    )[0]

    engine_args = stage_cfg["engine_args"]
    assert engine_args["distributed_executor_backend"] == "ray"
    assert engine_args["boundary_ratio"] == 0.875
    assert engine_args["flow_shift"] == 5.0
    assert engine_args["trust_remote_code"] is True


def test_default_stage_config_whitelist_none_fallback():
    """DeployConfig / StageDeployConfig whitelist fields with value None
    fall back to OmniDiffusionConfig dataclass defaults."""
    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
        {
            # DeployConfig pipeline-wide
            "trust_remote_code": None,
            "distributed_executor_backend": None,
            "dtype": None,
            # StageDeployConfig
            "enforce_eager": None,
        }
    )[0]

    engine_args = stage_cfg["engine_args"]

    assert engine_args["trust_remote_code"] is False
    assert engine_args["distributed_executor_backend"] == "mp"
    assert engine_args["dtype"] == "auto"
    assert engine_args["enforce_eager"] is False


def test_serve_cli_accepts_ulysses_mode():
    """Ensure diffusion serve CLI exposes ulysses_mode and wires it to parallel_config."""
    parser = TrackingArgumentParser()
    subparsers = parser.add_subparsers(dest="command")
    OmniServeCommand().subparser_init(subparsers)

    args = parser.parse_args(
        [
            "serve",
            "Qwen/Qwen-Image",
            "--omni",
            "--usp",
            "4",
            "--ulysses-mode",
            "advanced_uaa",
        ]
    )

    explicit_kwargs = args.get_explicit_kwargs_dict()
    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
    parallel_config = stage_cfg["engine_args"]["parallel_config"]

    assert args.ulysses_mode == "advanced_uaa"
    assert parallel_config.ulysses_degree == 4
    assert parallel_config.ulysses_mode == "advanced_uaa"


def test_serve_cli_accepts_diffusion_pipeline_profiler_flag():
    """Ensure diffusion serve CLI exposes the profiler switch."""
    parser = TrackingArgumentParser()
    subparsers = parser.add_subparsers(dest="command")
    OmniServeCommand().subparser_init(subparsers)

    args = parser.parse_args(
        [
            "serve",
            "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
            "--omni",
            "--enable-diffusion-pipeline-profiler",
        ]
    )

    explicit_kwargs = args.get_explicit_kwargs_dict()
    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]

    assert args.enable_diffusion_pipeline_profiler is True
    assert stage_cfg["engine_args"]["enable_diffusion_pipeline_profiler"] is True


def test_serve_cli_accepts_diffusion_attention_backend():
    """Ensure diffusion serve CLI exposes the shorthand backend flag."""
    parser = TrackingArgumentParser()
    subparsers = parser.add_subparsers(dest="command")
    OmniServeCommand().subparser_init(subparsers)

    args = parser.parse_args(
        [
            "serve",
            "Qwen/Qwen-Image",
            "--omni",
            "--diffusion-attention-backend",
            "FLASH_ATTN",
        ]
    )

    explicit_kwargs = args.get_explicit_kwargs_dict()
    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
    diffusion_attention_config = stage_cfg["engine_args"]["diffusion_attention_config"]

    assert args.diffusion_attention_backend == "FLASH_ATTN"
    assert isinstance(diffusion_attention_config, AttentionConfig)
    assert diffusion_attention_config.default is not None
    assert diffusion_attention_config.default.backend == "FLASH_ATTN"


def test_serve_cli_accepts_additional_config():
    """Ensure diffusion serve CLI exposes additional_config and forwards it to stage config."""
    parser = TrackingArgumentParser()
    subparsers = parser.add_subparsers(dest="command")
    OmniServeCommand().subparser_init(subparsers)

    args = parser.parse_args(
        [
            "serve",
            "Qwen/Qwen-Image",
            "--omni",
            "--additional-config",
            '{"torchair_graph_config":{"enabled":true}}',
        ]
    )

    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(vars(args))[0]

    engine_args = stage_cfg["engine_args"]

    assert args.additional_config == {"torchair_graph_config": {"enabled": True}}
    assert engine_args["additional_config"] == {"torchair_graph_config": {"enabled": True}}


def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mocker):
    """Ensure YAML/deploy stage resolution forwards top-level additional_config."""
    fake_diffusion_stage = SimpleNamespace(
        stage_type="diffusion",
        engine_args=SimpleNamespace(),
    )
    fake_llm_stage = SimpleNamespace(
        stage_type="llm",
        engine_args=SimpleNamespace(),
    )
    mocker.patch(
        "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs",
        return_value=("dummy.yaml", [fake_llm_stage, fake_diffusion_stage]),
    )

    engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
    engine._strip_single_engine_args = lambda kwargs: kwargs

    _, stage_configs = engine._resolve_stage_configs(
        "dummy-model",
        {
            "stage_configs_path": "dummy.yaml",
            "additional_config": {"torchair_graph_config": {"enabled": True}},
        },
    )

    assert not hasattr(stage_configs[0].engine_args, "additional_config")
    assert stage_configs[1].engine_args.additional_config == {"torchair_graph_config": {"enabled": True}}