"""Unit tests for PD (Prefill-Decode) disaggregation in the Omni orchestrator.

Tests the PD detection, validation, config parsing, sampling param
preparation, and routing logic added by the PD disaggregation feature
(issue #1188).  All tests run without GPU.

NOTE (v1908 adaptation): Tests that relied on the old OmniStage / stage_list
architecture (removed in PR #1908) are marked xfail with
``reason="Requires migration to v1908 Orchestrator architecture"``.
The remaining tests exercise PDDisaggregationMixin directly and work
without spinning up a real engine.
"""

import uuid
import warnings
from queue import Empty, Queue
from types import SimpleNamespace
from typing import Any

import pytest
from vllm import SamplingParams

from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin

pytestmark = pytest.mark.skip(reason="Temporarily skip PD entrypoint tests while PD config is being removed.")

# Suppress noisy DeprecationWarnings from optional Swig bindings imported by vLLM dependencies.
warnings.filterwarnings(
    "ignore",
    message=r"builtin type SwigPy.*has no __module__ attribute",
    category=DeprecationWarning,
)


def _ns(**kwargs):
    """Create a lightweight attribute object for tests."""
    return SimpleNamespace(**kwargs)


# ---------------------------------------------------------------------------
# Fake helpers (same pattern as test_omni_llm.py)
# ---------------------------------------------------------------------------


class _FakeEngineArgs(dict):
    """Fake engine args that supports both attribute and dict access."""

    def __init__(self, args_dict: dict[str, Any]):
        super().__init__(args_dict)
        if "model_stage" not in self:
            self["model_stage"] = None
        if "engine_output_type" not in self:
            self["engine_output_type"] = None
        for key, value in self.items():
            setattr(self, key, value)


class _FakeStageConfig:
    def __init__(self, config_dict: dict[str, Any]):
        engine_args_dict = config_dict.get("engine_args", {})
        self.engine_args = _FakeEngineArgs(engine_args_dict)
        self.final_output = config_dict.get("final_output", False)
        self.final_output_type = config_dict.get("final_output_type", None)
        self.stage_id = config_dict.get("stage_id", 0)
        self.is_prefill_only = config_dict.get("is_prefill_only", False)
        self.is_decode_only = config_dict.get("is_decode_only", False)
        self.engine_input_source = config_dict.get("engine_input_source", [])
        self.is_comprehension = config_dict.get("is_comprehension", False)
        self._config_dict = config_dict


class _FakeQueue:
    def __init__(self, maxsize=0):
        self._queue = Queue(maxsize=maxsize)

    def put(self, item):
        self._queue.put(item)

    def put_nowait(self, item):
        self._queue.put_nowait(item)

    def get(self):
        return self._queue.get()

    def get_nowait(self):
        return self._queue.get_nowait()

    def empty(self):
        return self._queue.empty()


class _FakeStage:
    """Lightweight stage stub with PD disaggregation flag support."""

    def __init__(self, config, stage_init_timeout: int = 300):
        if isinstance(config, dict):
            config = _FakeStageConfig(config)
        self.config = config
        self.stage_config = config
        self.engine = None
        self.engine_outputs = None
        self.stage_id = getattr(config, "stage_id", 0)
        self.engine_args = config.engine_args
        self.model_stage = getattr(config.engine_args, "model_stage", None)
        self.stage_type = "llm"
        self.default_sampling_params = SamplingParams(temperature=1.0)
        self.final_output = config.final_output if hasattr(config, "final_output") else False
        self.final_output_type = getattr(config, "final_output_type", None)
        self.is_prefill_only = getattr(config, "is_prefill_only", False)
        self.is_decode_only = getattr(config, "is_decode_only", False)
        self.engine_input_source = getattr(config, "engine_input_source", [])
        self.is_comprehension = getattr(config, "is_comprehension", False)
        processed_input = getattr(config, "_config_dict", {}).get("processed_input", ["processed"])
        self._processed_input = processed_input
        self._in_q = None
        self._out_q = None
        self._proc = None
        self._stage_init_timeout = max(0, int(stage_init_timeout))

    def attach_queues(self, in_q, out_q):
        self._in_q = in_q
        self._out_q = out_q

    def init_stage_worker(
        self, model: str, *, is_async=False, shm_threshold_bytes=65536, ctx=None, batch_timeout=10, **kwargs
    ):
        self._proc = _ns(
            start=lambda: None,
            join=lambda timeout=None: None,
            is_alive=lambda: False,
            terminate=lambda: None,
        )
        if self._out_q is not None:
            try:
                self._out_q.put_nowait({"type": "stage_ready", "stage_id": self.stage_id})
            except Exception:
                pass

    def stop_stage_worker(self):
        if self._in_q is not None:
            try:
                self._in_q.put_nowait({"type": "shutdown"})
            except Exception:
                pass

    def submit(self, payload: dict[str, Any]):
        if self._in_q is not None:
            self._in_q.put(payload)

    def try_collect(self) -> Any:
        if self._out_q is None:
            return None
        try:
            return self._out_q.get_nowait()
        except Empty:
            return None

    def set_engine_outputs(self, outputs):
        self.engine_outputs = outputs

    def process_engine_inputs(self, stage_list, prompts):
        return self._processed_input


# ---------------------------------------------------------------------------
# Shared mock setup helpers
# ---------------------------------------------------------------------------


def _setup_engine_mocks(monkeypatch):
    fake_engine = _ns()
    fake_engine.tokenizer = _ns()
    fake_engine.log_stats = False
    fake_engine.vllm_config = _ns()
    fake_engine.vllm_config.model_config = _ns()
    fake_engine.vllm_config.model_config.io_processor_plugin = None
    fake_engine.get_supported_tasks = lambda: []
    fake_engine.model_config = _ns()
    fake_engine.model_config.io_processor_plugin = None
    fake_registry = _ns()
    fake_registry.resolve_model_cls = lambda *args, **kwargs: (_ns(), "test_arch")
    fake_engine.model_config.registry = fake_registry
    fake_engine.vllm_config.model_config.registry = fake_registry

    monkeypatch.setattr(
        "vllm.v1.engine.llm_engine.LLMEngine.from_engine_args",
        lambda **kw: fake_engine,
        raising=False,
    )

    class FakeModelClass:
        pass

    monkeypatch.setattr(
        "vllm.model_executor.model_loader.utils.get_model_architecture",
        lambda model_config: (FakeModelClass, "test_arch"),
        raising=False,
    )
    monkeypatch.setattr(
        "vllm.model_executor.model_loader.utils._get_model_architecture",
        lambda model_config: (FakeModelClass, "test_arch"),
        raising=False,
    )
    monkeypatch.setattr(
        "vllm.model_executor.models.adapters.try_create_mm_pooling_model_cls",
        lambda model_cls: model_cls,
        raising=False,
    )
    monkeypatch.setattr(
        "vllm.multimodal.cache._enable_processor_cache",
        lambda model_config, mm_registry: False,
        raising=False,
    )
    monkeypatch.setattr(
        "vllm.plugins.io_processors.get_io_processor",
        lambda vllm_config, io_processor_plugin: None,
        raising=False,
    )


def _setup_multiprocessing_mocks(monkeypatch):
    import multiprocessing as mp

    fake_process_instance = _ns(
        start=lambda: None,
        join=lambda timeout=None: None,
        is_alive=lambda: False,
        terminate=lambda: None,
    )

    def fake_process_class(*args, **kwargs):
        return fake_process_instance

    fake_ctx = _ns()
    fake_ctx.Queue = lambda maxsize=0: _FakeQueue(maxsize=maxsize)
    fake_ctx.Process = fake_process_class

    monkeypatch.setattr(mp, "get_context", lambda method: fake_ctx, raising=False)
    monkeypatch.setattr(mp, "Process", fake_process_class, raising=False)


def _setup_ipc_mocks(monkeypatch):
    # These IPC helpers existed in the old architecture; no-op in new arch.
    pass


def _setup_log_mocks(monkeypatch):
    class _FakeOrchestratorAggregator:
        def __init__(self, num_stages, enable_stats, wall_start_ts, final_stage_id_for_e2e=None):
            self.num_stages = num_stages
            self.enable_stats = enable_stats
            self.stage_first_ts = [None] * num_stages
            self.stage_last_ts = [None] * num_stages
            self.stage_total_tokens = [0] * num_stages
            self.accumulated_gen_time_ms = {}
            self.e2e_done = set()
            self.e2e_count = 0
            self.e2e_total_ms = 0.0

        def on_stage_metrics(self, stage_id, req_id, metrics, final_output_type=None):
            pass

        def on_finalize_request(self, stage_id, req_id, start_ts):
            self.e2e_done.add(req_id)

        def on_forward(self, from_stage, to_stage, req_id, size_bytes, tx_ms, use_shm):
            pass

        def accumulate_diffusion_metrics(self, stage_type, req_id, engine_outputs):
            pass

        def record_audio_generated_frames(self, output, stage_id, req_id):
            pass

        def stage_postprocess_timer(self, stage_id, req_id):
            from contextlib import contextmanager

            @contextmanager
            def _noop():
                yield

            return _noop()

        def build_and_log_summary(self):
            return "Fake summary"

    monkeypatch.setattr(
        "vllm_omni.entrypoints.omni.OrchestratorAggregator",
        _FakeOrchestratorAggregator,
        raising=False,
    )


def _clear_modules():
    import sys

    for module_name in [
        "vllm_omni.entrypoints.utils",
        "vllm_omni.entrypoints.omni",
    ]:
        if module_name in sys.modules:
            del sys.modules[module_name]


@pytest.fixture(autouse=True)
def mock_get_config(monkeypatch):
    """Auto-mock get_config and related model loading functions."""
    import sys

    fake_tokenizer = _ns()
    fake_tokenizer.encode = lambda *args, **kwargs: [1, 2, 3]
    fake_tokenizer.decode = lambda *args, **kwargs: "test"

    def _mock_init_tokenizer_from_configs(model_config=None, **kwargs):
        return fake_tokenizer

    monkeypatch.setattr(
        "vllm.transformers_utils.tokenizer.init_tokenizer_from_configs",
        _mock_init_tokenizer_from_configs,
        raising=False,
    )
    tokenizer_module_path = "vllm.transformers_utils.tokenizer"
    if tokenizer_module_path in sys.modules:
        setattr(sys.modules[tokenizer_module_path], "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs)

    def _mock_length_from_prompt_token_ids_or_embeds(prompt_token_ids=None, prompt_embeds=None):
        if prompt_token_ids is not None:
            if isinstance(prompt_token_ids, list):
                return len(prompt_token_ids)
        return 10

    monkeypatch.setattr(
        "vllm.utils.length_from_prompt_token_ids_or_embeds", _mock_length_from_prompt_token_ids_or_embeds, raising=False
    )
    monkeypatch.setattr(
        "vllm_omni.engine.input_processor.length_from_prompt_token_ids_or_embeds",
        _mock_length_from_prompt_token_ids_or_embeds,
        raising=False,
    )

    processor_module_path = "vllm_omni.engine.input_processor"
    if processor_module_path in sys.modules:
        setattr(
            sys.modules[processor_module_path],
            "length_from_prompt_token_ids_or_embeds",
            _mock_length_from_prompt_token_ids_or_embeds,
        )

    monkeypatch.setattr(
        "vllm_omni.entrypoints.async_omni.init_tokenizer_from_configs", _mock_init_tokenizer_from_configs, raising=False
    )
    async_omni_path = "vllm_omni.entrypoints.async_omni"
    if async_omni_path in sys.modules:
        setattr(sys.modules[async_omni_path], "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs)

    fake_hf_config = _ns()
    fake_hf_config.model_type = "qwen2_5_omni"

    monkeypatch.setattr(
        "vllm.transformers_utils.config.get_config", lambda model, **kwargs: fake_hf_config, raising=False
    )
    monkeypatch.setattr("vllm_omni.entrypoints.utils.get_config", lambda model, **kwargs: fake_hf_config, raising=False)

    def _mock_cached_file(path_or_repo_id, *args, **kwargs):
        import os
        import tempfile

        fake_config_file = os.path.join(tempfile.gettempdir(), "fake_config.json")
        if not os.path.exists(fake_config_file):
            with open(fake_config_file, "w") as f:
                f.write('{"model_type": "qwen2_5_omni"}')
        return fake_config_file

    monkeypatch.setattr("transformers.utils.hub.cached_file", _mock_cached_file, raising=False)
    monkeypatch.setattr(
        "transformers.utils.hub.cached_files",
        lambda path_or_repo_id, filenames, **kwargs: (
            [_mock_cached_file(path_or_repo_id, filenames[0])] if filenames else None
        ),
        raising=False,
    )


# ---------------------------------------------------------------------------
# Helper to build an Omni instance with PD stage configs
# ---------------------------------------------------------------------------


def _make_pd_omni(monkeypatch, stage_configs, *, extra_setup=None):
    """Create a lightweight PDDisaggregationMixin instance for unit tests.

    Bypasses the full OmniBase / AsyncOmniEngine init chain so tests run
    without GPU.  Returns an object that has all PDDisaggregationMixin
    methods and state (``_pd_separation_pair``, ``_pd_kv_params_by_req``,
    etc.) initialised from *stage_configs*.

    Tests that need the full ``Omni.generate()`` loop (old stage_list / queue
    infrastructure) are marked ``xfail`` and not covered here.
    """
    configs = [_FakeStageConfig(c) for c in stage_configs]

    class _LightweightOmni(PDDisaggregationMixin):
        """Minimal shim: exposes stage_configs so PDDisaggregationMixin works."""

        def __init__(self):
            self._name = "Omni"
            self._stage_configs = configs
            self._init_pd_state()

        @property
        def stage_configs(self):
            return self._stage_configs

    if extra_setup:
        import vllm_omni.entrypoints.omni as omni_module

        extra_setup(monkeypatch, omni_module)

    return _LightweightOmni()


# ---------------------------------------------------------------------------
# Stage config templates
# ---------------------------------------------------------------------------


def _prefill_stage_cfg(stage_id=0, **overrides):
    cfg = {
        "stage_id": stage_id,
        "engine_args": {
            "model_stage": "thinker",
            "kv_transfer_config": {
                "kv_connector": "MooncakeConnector",
                "kv_role": "kv_producer",
                "kv_rank": 0,
                "kv_parallel_size": 2,
                "kv_connector_extra_config": {"mooncake_bootstrap_port": 25201},
            },
        },
        "is_prefill_only": True,
        "final_output": False,
        "is_comprehension": True,
    }
    cfg.update(overrides)
    return cfg


def _decode_stage_cfg(stage_id=1, engine_input_source=None, **overrides):
    cfg = {
        "stage_id": stage_id,
        "engine_args": {
            "model_stage": "thinker",
            "kv_transfer_config": {
                "kv_connector": "MooncakeConnector",
                "kv_role": "kv_consumer",
                "kv_rank": 1,
                "kv_parallel_size": 2,
                "kv_connector_extra_config": {"mooncake_bootstrap_port": 25202},
            },
        },
        "is_decode_only": True,
        "engine_input_source": engine_input_source if engine_input_source is not None else [0],
        "final_output": True,
        "final_output_type": "text",
        "is_comprehension": True,
    }
    cfg.update(overrides)
    return cfg


def _talker_stage_cfg(stage_id=2, engine_input_source=None, **overrides):
    cfg = {
        "stage_id": stage_id,
        "engine_args": {"model_stage": "talker"},
        "engine_input_source": engine_input_source if engine_input_source is not None else [1],
        "final_output": False,
    }
    cfg.update(overrides)
    return cfg


def _code2wav_stage_cfg(stage_id=3, engine_input_source=None, **overrides):
    cfg = {
        "stage_id": stage_id,
        "engine_args": {"model_stage": "code2wav"},
        "engine_input_source": engine_input_source if engine_input_source is not None else [2],
        "final_output": True,
        "final_output_type": "audio",
    }
    cfg.update(overrides)
    return cfg


# ===================================================================
# Tests: PD pair detection
# ===================================================================


class TestDetectPDSeparation:
    """Tests for Omni._detect_pd_separation()."""

    def test_detects_pd_pair(self, monkeypatch):
        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(stage_id=0),
                _decode_stage_cfg(stage_id=1, engine_input_source=[0]),
            ],
        )
        assert omni._pd_separation_pair == (0, 1)

    def test_no_pd_pair_without_flags(self, monkeypatch):
        """Normal (non-PD) pipeline has no PD pair."""
        omni = _make_pd_omni(
            monkeypatch,
            [
                {
                    "stage_id": 0,
                    "engine_args": {"model_stage": "thinker"},
                    "final_output": True,
                    "final_output_type": "text",
                },
                {
                    "stage_id": 1,
                    "engine_args": {"model_stage": "talker"},
                    "engine_input_source": [0],
                    "final_output": True,
                    "final_output_type": "audio",
                },
            ],
        )
        assert omni._pd_separation_pair is None

    def test_detects_pd_pair_in_4_stage_pipeline(self, monkeypatch):
        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(stage_id=0),
                _decode_stage_cfg(stage_id=1, engine_input_source=[0]),
                _talker_stage_cfg(stage_id=2, engine_input_source=[1]),
                _code2wav_stage_cfg(stage_id=3, engine_input_source=[2]),
            ],
        )
        assert omni._pd_separation_pair == (0, 1)

    def test_pd_pair_uses_stage_id_for_input_source(self, monkeypatch):
        """engine_input_source references stage_id, not list index."""
        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(stage_id=10),
                _decode_stage_cfg(stage_id=20, engine_input_source=[10]),
            ],
        )
        assert omni._pd_separation_pair == (0, 1)


# ===================================================================
# Tests: PD config validation
# ===================================================================


class TestValidatePDConfig:
    """Tests for Omni._validate_pd_separation_config()."""

    def test_valid_config_passes(self, monkeypatch):
        """Valid PD config should not raise."""
        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(),
                _decode_stage_cfg(engine_input_source=[0]),
            ],
        )
        # If we got here without error, validation passed
        assert omni._pd_separation_pair == (0, 1)

    def test_mismatched_connector_raises(self, monkeypatch):
        """Different kv_connector types should raise ValueError."""
        decode_cfg = _decode_stage_cfg(engine_input_source=[0])
        decode_cfg["engine_args"]["kv_transfer_config"]["kv_connector"] = "NixlConnector"

        with pytest.raises(ValueError, match="connector mismatch"):
            _make_pd_omni(monkeypatch, [_prefill_stage_cfg(), decode_cfg])

    def test_wrong_prefill_role_raises(self, monkeypatch):
        """Prefill with kv_consumer role should raise."""
        prefill_cfg = _prefill_stage_cfg()
        prefill_cfg["engine_args"]["kv_transfer_config"]["kv_role"] = "kv_consumer"

        with pytest.raises(ValueError, match="kv_role must be"):
            _make_pd_omni(monkeypatch, [prefill_cfg, _decode_stage_cfg(engine_input_source=[0])])

    def test_wrong_decode_role_raises(self, monkeypatch):
        """Decode with kv_producer role should raise."""
        decode_cfg = _decode_stage_cfg(engine_input_source=[0])
        decode_cfg["engine_args"]["kv_transfer_config"]["kv_role"] = "kv_producer"

        with pytest.raises(ValueError, match="kv_role must be"):
            _make_pd_omni(monkeypatch, [_prefill_stage_cfg(), decode_cfg])

    def test_missing_kv_transfer_config_raises(self, monkeypatch):
        """Missing kv_transfer_config should raise."""
        prefill_cfg = _prefill_stage_cfg()
        del prefill_cfg["engine_args"]["kv_transfer_config"]

        with pytest.raises(ValueError, match="kv_transfer_config"):
            _make_pd_omni(monkeypatch, [prefill_cfg, _decode_stage_cfg(engine_input_source=[0])])

    def test_mismatched_buffer_device_raises(self, monkeypatch):
        """Mismatched kv_buffer_device should raise."""
        prefill_cfg = _prefill_stage_cfg()
        prefill_cfg["engine_args"]["kv_transfer_config"]["kv_buffer_device"] = "cuda"
        decode_cfg = _decode_stage_cfg(engine_input_source=[0])
        decode_cfg["engine_args"]["kv_transfer_config"]["kv_buffer_device"] = "cpu"

        with pytest.raises(ValueError, match="kv_buffer_device mismatch"):
            _make_pd_omni(monkeypatch, [prefill_cfg, decode_cfg])


# ===================================================================
# Tests: Connector info extraction
# ===================================================================


class TestGetPDConnectorInfo:
    """Tests for Omni._get_pd_connector_info()."""

    def test_extracts_bootstrap_addr_for_mooncake(self, monkeypatch):
        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(),
                _decode_stage_cfg(engine_input_source=[0]),
            ],
        )
        info = omni._pd_connector_info
        assert "prefill_bootstrap_addr" in info
        assert info["prefill_bootstrap_addr"] == "127.0.0.1:25201"

    def test_none_for_non_pd_pipeline(self, monkeypatch):
        omni = _make_pd_omni(
            monkeypatch,
            [
                {"stage_id": 0, "engine_args": {}, "final_output": True, "final_output_type": "text"},
            ],
        )
        assert omni._pd_connector_info is None


# ===================================================================
# Tests: Prefill sampling params preparation
# ===================================================================


class TestPreparePrefillSamplingParams:
    """Tests for Omni._prepare_prefill_sampling_params()."""

    def test_sets_max_tokens_to_1(self, monkeypatch):
        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(),
                _decode_stage_cfg(engine_input_source=[0]),
            ],
        )
        sp = SamplingParams(max_tokens=2048)
        result = omni._prepare_prefill_sampling_params("req-1", sp)

        assert result.max_tokens == 1
        assert result is not sp  # should be cloned

    def test_injects_kv_transfer_params(self, monkeypatch):
        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(),
                _decode_stage_cfg(engine_input_source=[0]),
            ],
        )
        sp = SamplingParams(max_tokens=2048)
        result = omni._prepare_prefill_sampling_params("req-1", sp)

        kv_params = result.extra_args["kv_transfer_params"]
        assert kv_params["do_remote_decode"] is True
        assert kv_params["do_remote_prefill"] is False
        assert kv_params["transfer_id"] == "xfer-req-1"

    def test_preserves_existing_extra_args(self, monkeypatch):
        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(),
                _decode_stage_cfg(engine_input_source=[0]),
            ],
        )
        sp = SamplingParams(max_tokens=2048, extra_args={"custom_key": "value"})
        result = omni._prepare_prefill_sampling_params("req-1", sp)

        assert result.extra_args["custom_key"] == "value"
        assert "kv_transfer_params" in result.extra_args

    def test_does_not_mutate_original(self, monkeypatch):
        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(),
                _decode_stage_cfg(engine_input_source=[0]),
            ],
        )
        sp = SamplingParams(max_tokens=2048)
        _ = omni._prepare_prefill_sampling_params("req-1", sp)

        assert sp.max_tokens == 2048
        assert sp.extra_args is None


# ===================================================================
# Tests: Sampling params auto-duplication for PD split
# ===================================================================


@pytest.mark.xfail(reason="Requires migration to v1908 Orchestrator architecture (no stage_list / OmniStage)")
class TestSamplingParamsAutoDuplication:
    """When user provides N-1 sampling params (for logical stages), the
    orchestrator should auto-duplicate the thinker params for the decode stage.
    """

    def test_auto_duplicates_for_4_stage_pipeline(self, monkeypatch):
        """User provides 3 params for 4 physical stages -> auto-insert decode params."""
        test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000001")

        def _extra_setup(mp, omni_module):
            mp.setattr(uuid, "uuid4", lambda: test_uuid)
            mp.setattr(omni_module, "uuid", uuid)

        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(stage_id=0),
                _decode_stage_cfg(stage_id=1, engine_input_source=[0]),
                _talker_stage_cfg(stage_id=2, engine_input_source=[1]),
                _code2wav_stage_cfg(stage_id=3, engine_input_source=[2]),
            ],
            extra_setup=_extra_setup,
        )

        assert omni._pd_separation_pair == (0, 1)
        assert len(omni.stage_list) == 4

        # Simulate outputs for all stages
        expected_rid = f"0_{test_uuid}"
        for i in range(4):
            omni.stage_list[i]._out_q.put_nowait(
                {
                    "request_id": expected_rid,
                    "engine_outputs": [_ns(request_id=expected_rid, outputs=[_ns(token_ids=[1, 2])])],
                    "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
                }
            )

        # Provide 3 params (one less than 4 stages) - should auto-duplicate
        sp_thinker = SamplingParams(temperature=0.4, max_tokens=2048)
        sp_talker = SamplingParams(temperature=0.9, max_tokens=4096)
        sp_code2wav = SamplingParams(temperature=0.0, max_tokens=65536)

        # This should NOT raise ValueError about param count mismatch
        outputs = omni.generate(
            prompts=["hello"],
            sampling_params_list=[sp_thinker, sp_talker, sp_code2wav],
        )
        assert isinstance(outputs, list)


# ===================================================================
# Tests: KV transfer params normalization
# ===================================================================


class TestNormalizeKVTransferParams:
    def test_dict_passthrough(self, monkeypatch):
        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(),
                _decode_stage_cfg(engine_input_source=[0]),
            ],
        )
        d = {"transfer_id": "test", "do_remote_decode": True}
        assert omni._normalize_kv_transfer_params(d) is d

    def test_none_returns_none(self, monkeypatch):
        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(),
                _decode_stage_cfg(engine_input_source=[0]),
            ],
        )
        assert omni._normalize_kv_transfer_params(None) is None

    def test_dataclass_to_dict(self, monkeypatch):
        from dataclasses import dataclass

        @dataclass
        class FakeKVParams:
            transfer_id: str = "test"
            do_remote_decode: bool = True

        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(),
                _decode_stage_cfg(engine_input_source=[0]),
            ],
        )
        result = omni._normalize_kv_transfer_params(FakeKVParams())
        assert isinstance(result, dict)
        assert result["transfer_id"] == "test"


# ===================================================================
# Tests: _kv_cfg_to_dict
# ===================================================================


class TestKvCfgToDict:
    def test_dict_passthrough(self, monkeypatch):
        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(),
                _decode_stage_cfg(engine_input_source=[0]),
            ],
        )
        d = {"kv_connector": "MooncakeConnector"}
        assert omni._kv_cfg_to_dict(d) is d

    def test_none_returns_empty(self, monkeypatch):
        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(),
                _decode_stage_cfg(engine_input_source=[0]),
            ],
        )
        assert omni._kv_cfg_to_dict(None) == {}

    def test_dataclass_converted(self, monkeypatch):
        from dataclasses import dataclass

        @dataclass
        class FakeCfg:
            kv_connector: str = "TestConnector"
            kv_role: str = "kv_producer"

        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(),
                _decode_stage_cfg(engine_input_source=[0]),
            ],
        )
        result = omni._kv_cfg_to_dict(FakeCfg())
        assert result["kv_connector"] == "TestConnector"
        assert result["kv_role"] == "kv_producer"


# ===================================================================
# Tests: PD routing in scheduling loop
# ===================================================================


@pytest.mark.xfail(reason="Requires migration to v1908 Orchestrator architecture (no stage_list / OmniStage)")
class TestPDRouting:
    """Test that the scheduling loop correctly routes requests from
    prefill to decode stage with proper kv_transfer_params.
    """

    def test_prefill_stage_receives_max_tokens_1(self, monkeypatch):
        """Stage 0 (prefill) should receive max_tokens=1."""
        test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000002")

        def _extra_setup(mp, omni_module):
            mp.setattr(uuid, "uuid4", lambda: test_uuid)
            mp.setattr(omni_module, "uuid", uuid)

        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(stage_id=0),
                _decode_stage_cfg(stage_id=1, engine_input_source=[0]),
            ],
            extra_setup=_extra_setup,
        )

        expected_rid = f"0_{test_uuid}"

        # Put stage outputs in both queues
        omni.stage_list[0]._out_q.put_nowait(
            {
                "request_id": expected_rid,
                "engine_outputs": [_ns(request_id=expected_rid, outputs=[_ns(token_ids=[1])])],
                "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
            }
        )
        omni.stage_list[1]._out_q.put_nowait(
            {
                "request_id": expected_rid,
                "engine_outputs": [_ns(request_id=expected_rid, outputs=[_ns(token_ids=[1, 2, 3])])],
                "metrics": {"num_tokens_out": 3, "stage_gen_time_ms": 50.0},
            }
        )

        sp_list = [SamplingParams(max_tokens=2048), SamplingParams(max_tokens=2048)]
        omni.generate(prompts=["hello"], sampling_params_list=sp_list)

        # Check what was submitted to stage 0's input queue
        # (skip the stage_ready message first)
        task = omni.stage_list[0]._in_q.get_nowait()
        assert task["sampling_params"].max_tokens == 1
        kv_params = task["sampling_params"].extra_args["kv_transfer_params"]
        assert kv_params["do_remote_decode"] is True
        assert kv_params["do_remote_prefill"] is False
        assert kv_params["transfer_id"] == f"xfer-{expected_rid}"

    def test_decode_stage_receives_original_prompt(self, monkeypatch):
        """Decode stage should get the original prompt (not processed outputs)."""
        test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000003")

        def _extra_setup(mp, omni_module):
            mp.setattr(uuid, "uuid4", lambda: test_uuid)
            mp.setattr(omni_module, "uuid", uuid)

        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(stage_id=0),
                _decode_stage_cfg(stage_id=1, engine_input_source=[0]),
            ],
            extra_setup=_extra_setup,
        )

        expected_rid = f"0_{test_uuid}"
        original_prompt = "test prompt for PD"

        omni.stage_list[0]._out_q.put_nowait(
            {
                "request_id": expected_rid,
                "engine_outputs": [_ns(request_id=expected_rid, outputs=[_ns(token_ids=[1])])],
                "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
            }
        )
        omni.stage_list[1]._out_q.put_nowait(
            {
                "request_id": expected_rid,
                "engine_outputs": [_ns(request_id=expected_rid, outputs=[_ns(token_ids=[1, 2, 3])])],
                "metrics": {"num_tokens_out": 3, "stage_gen_time_ms": 50.0},
            }
        )

        sp_list = [SamplingParams(max_tokens=2048), SamplingParams(max_tokens=2048)]
        omni.generate(prompts=[original_prompt], sampling_params_list=sp_list)

        # Check what was forwarded to stage 1 (decode)
        # The connector sends tasks to stage 1's input queue
        task = omni.stage_list[1]._in_q.get_nowait()
        # The engine_inputs should contain the original prompt
        engine_inputs = task.get("engine_inputs")
        # For PD routing, the original prompt is wrapped in a list
        if isinstance(engine_inputs, list):
            assert original_prompt in engine_inputs
        else:
            assert engine_inputs == original_prompt

    def test_decode_kv_params_have_correct_flags(self, monkeypatch):
        """Decode stage kv_transfer_params should have correct role flags."""
        test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000004")

        def _extra_setup(mp, omni_module):
            mp.setattr(uuid, "uuid4", lambda: test_uuid)
            mp.setattr(omni_module, "uuid", uuid)

        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(stage_id=0),
                _decode_stage_cfg(stage_id=1, engine_input_source=[0]),
            ],
            extra_setup=_extra_setup,
        )

        expected_rid = f"0_{test_uuid}"

        omni.stage_list[0]._out_q.put_nowait(
            {
                "request_id": expected_rid,
                "engine_outputs": [_ns(request_id=expected_rid, outputs=[_ns(token_ids=[1])])],
                "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
            }
        )
        omni.stage_list[1]._out_q.put_nowait(
            {
                "request_id": expected_rid,
                "engine_outputs": [_ns(request_id=expected_rid, outputs=[_ns(token_ids=[1, 2, 3])])],
                "metrics": {"num_tokens_out": 3, "stage_gen_time_ms": 50.0},
            }
        )

        sp_list = [SamplingParams(max_tokens=2048), SamplingParams(max_tokens=2048)]
        omni.generate(prompts=["hello"], sampling_params_list=sp_list)

        # Check decode task's kv_transfer_params
        task = omni.stage_list[1]._in_q.get_nowait()
        kv_params = task["sampling_params"].extra_args["kv_transfer_params"]
        assert kv_params["do_remote_prefill"] is True
        assert kv_params["do_remote_decode"] is False
        assert kv_params["transfer_id"] == f"xfer-{expected_rid}"
        assert kv_params["remote_bootstrap_addr"] == "127.0.0.1:25201"


# ===================================================================
# Tests: KV params cleanup
# ===================================================================


class TestKVParamsCleanup:
    def test_drop_cleans_up(self, monkeypatch):
        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(),
                _decode_stage_cfg(engine_input_source=[0]),
            ],
        )
        omni._pd_kv_params_by_req["req-1"] = {"transfer_id": "xfer-1"}
        omni._drop_pd_kv_params("req-1")
        assert "req-1" not in omni._pd_kv_params_by_req

    def test_drop_nonexistent_is_noop(self, monkeypatch):
        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(),
                _decode_stage_cfg(engine_input_source=[0]),
            ],
        )
        omni._drop_pd_kv_params("nonexistent")  # should not raise

    def test_pop_returns_stored_params(self, monkeypatch):
        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(),
                _decode_stage_cfg(engine_input_source=[0]),
            ],
        )
        stored = {"transfer_id": "xfer-1", "extra_field": "value"}
        omni._pd_kv_params_by_req["req-1"] = stored

        result = omni._pop_pd_kv_params("req-1")
        assert result == stored
        assert "req-1" not in omni._pd_kv_params_by_req

    def test_pop_uses_fallback_when_no_stored(self, monkeypatch):
        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(),
                _decode_stage_cfg(engine_input_source=[0]),
            ],
        )
        fallback = {"transfer_id": "xfer-fallback"}
        result = omni._pop_pd_kv_params("req-1", fallback=fallback)
        assert result == fallback


# ===================================================================
# Tests: Config YAML loads without error
# ===================================================================


class TestPDYAMLConfig:
    def test_pd_yaml_loads(self):
        """The PD separation YAML config should load without errors."""
        import os

        yaml_path = os.path.join(
            os.path.dirname(__file__),
            "../../vllm_omni/model_executor/stage_configs/qwen3_omni_moe_pd_separation.yaml",
        )
        yaml_path = os.path.abspath(yaml_path)
        if not os.path.exists(yaml_path):
            pytest.skip("PD separation YAML not found")

        from omegaconf import OmegaConf

        cfg = OmegaConf.load(yaml_path)
        stages = cfg.stage_args
        assert len(stages) == 4

        # Prefill stage
        assert stages[0].is_prefill_only is True
        assert stages[0].final_output is False
        assert stages[0].is_comprehension is True

        # Decode stage
        assert stages[1].is_decode_only is True
        assert stages[1].final_output is True
        assert stages[1].final_output_type == "text"
        assert stages[1].is_comprehension is True
        assert 0 in stages[1].engine_input_source

        # KV transfer configs
        assert stages[0].engine_args.kv_transfer_config.kv_role == "kv_producer"
        assert stages[1].engine_args.kv_transfer_config.kv_role == "kv_consumer"
        assert stages[0].engine_args.kv_transfer_config.kv_connector == "MooncakeConnector"
        assert stages[1].engine_args.kv_transfer_config.kv_connector == "MooncakeConnector"


class TestPrefillStopNeutralization:
    """Tests that _prepare_prefill_sampling_params neutralizes stop
    conditions to ensure finish_reason='length'.
    """

    def test_clears_stop_strings(self, monkeypatch):
        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(),
                _decode_stage_cfg(engine_input_source=[0]),
            ],
        )
        sp = SamplingParams(max_tokens=2048, stop=["</s>", "STOP"])
        result = omni._prepare_prefill_sampling_params("req-1", sp)
        assert result.stop == []

    def test_clears_stop_token_ids(self, monkeypatch):
        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(),
                _decode_stage_cfg(engine_input_source=[0]),
            ],
        )
        sp = SamplingParams(max_tokens=2048, stop_token_ids=[151643, 151644])
        result = omni._prepare_prefill_sampling_params("req-1", sp)
        assert result.stop_token_ids == []

    def test_clears_include_stop_str_in_output(self, monkeypatch):
        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(),
                _decode_stage_cfg(engine_input_source=[0]),
            ],
        )
        sp = SamplingParams(max_tokens=2048, include_stop_str_in_output=True)
        result = omni._prepare_prefill_sampling_params("req-1", sp)
        assert result.include_stop_str_in_output is False

    def test_original_sp_unchanged(self, monkeypatch):
        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(),
                _decode_stage_cfg(engine_input_source=[0]),
            ],
        )
        sp = SamplingParams(max_tokens=2048, stop=["</s>"], stop_token_ids=[151643])
        _ = omni._prepare_prefill_sampling_params("req-1", sp)
        assert sp.stop == ["</s>"]
        assert sp.stop_token_ids == [151643]


# ===================================================================
# Tests: Failure mode & memory leak prevention
# ===================================================================
# NOTE: Full generate()-level failure mode tests are removed for now.
# The _run_generation error handler (line 1344-1350 in omni.py) calls
# _drop_pd_kv_params but does not increment completed_requests, causing
# the while-loop to hang.  These tests need to be revisited once the
# production error-handling path is fixed to properly terminate on
# stage errors.


# ===================================================================
# Tests: TP size validation
# ===================================================================


class TestTPSizeValidation:
    """Tests that _validate_pd_separation_config checks tensor_parallel_size."""

    def test_matching_tp_passes(self, monkeypatch):
        """Same TP size should not raise."""
        prefill_cfg = _prefill_stage_cfg()
        prefill_cfg["engine_args"]["tensor_parallel_size"] = 2
        decode_cfg = _decode_stage_cfg(engine_input_source=[0])
        decode_cfg["engine_args"]["tensor_parallel_size"] = 2
        omni = _make_pd_omni(monkeypatch, [prefill_cfg, decode_cfg])
        assert omni._pd_separation_pair == (0, 1)

    def test_mismatched_tp_raises(self, monkeypatch):
        """Different TP sizes should raise ValueError."""
        prefill_cfg = _prefill_stage_cfg()
        prefill_cfg["engine_args"]["tensor_parallel_size"] = 2
        decode_cfg = _decode_stage_cfg(engine_input_source=[0])
        decode_cfg["engine_args"]["tensor_parallel_size"] = 4
        with pytest.raises(ValueError, match="tensor_parallel_size"):
            _make_pd_omni(monkeypatch, [prefill_cfg, decode_cfg])

    def test_default_tp_no_error(self, monkeypatch):
        """Stages without explicit TP (defaults to 1) should pass."""
        omni = _make_pd_omni(
            monkeypatch,
            [
                _prefill_stage_cfg(),
                _decode_stage_cfg(engine_input_source=[0]),
            ],
        )
        assert omni._pd_separation_pair == (0, 1)