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

"""MRV2 admission, capture, dispatch and request lifecycle contracts."""

from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import pytest
import torch
from vllm import SamplingParams
from vllm.config.compilation import CUDAGraphMode
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.worker.gpu.model_runner import GPUModelRunner
from vllm.v1.worker.gpu.sample.sampler import Sampler

from vllm_omni.model_executor.models.output_templates import OmniOutput
from vllm_omni.worker_v2.model_states import init_omni_model_state
from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner

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


def _make_runner():
    """Create an OmniGPUModelRunner without calling __init__."""
    runner = object.__new__(OmniGPUModelRunner)
    runner.model = MagicMock()
    runner.req_states = SimpleNamespace(req_id_to_index={"r1": 0, "r2": 1})
    runner.execute_model_state = None
    return runner


def test_add_requests_empty_admission_and_stop_id_sanitization():
    runner = _make_runner()
    runner.sampler = Sampler.__new__(Sampler)
    with patch.object(GPUModelRunner, "add_requests", return_value=None) as parent:
        runner.add_requests(SchedulerOutput.make_empty())
        parent.assert_not_called()

        # New requests always go upstream; narrow logits heads sanitize stop ids.
        sampling_params = SamplingParams(min_tokens=2, stop_token_ids=[2150])
        sampling_params.update_from_generation_config({}, 151645)
        runner.model = SimpleNamespace(logits_processor=SimpleNamespace(vocab_size=3072))
        output = SchedulerOutput.make_empty()
        output.scheduled_new_reqs = [SimpleNamespace(sampling_params=sampling_params)]
        runner.add_requests(output)
        parent.assert_called_once_with(output)
    assert sampling_params.all_stop_token_ids == {2150}
    assert sampling_params.eos_token_id == 151645


def test_prepare_native_data_plane_terminal_abort_split_and_warmup_skip():
    runner = _make_runner()
    plane = SimpleNamespace(
        register_request=MagicMock(),
        register_receivers=MagicMock(),
        request_terminal=MagicMock(),
        abort_requests=MagicMock(),
    )
    runner._omni_data_plane = plane
    new_req = SimpleNamespace(req_id="r1")
    warmup_req = SimpleNamespace(req_id="_warmup_0_")
    handle = SimpleNamespace(request_id="r2", external_req_id="ext-r2")
    scheduler_output = SimpleNamespace(
        scheduled_new_reqs=[new_req, warmup_req],
        pending_input_registrations=[handle],
        data_plane_terminal_req_ids={"r0"},
        finished_req_ids={"r0", "aborted"},
    )

    runner._prepare_native_data_plane(scheduler_output)

    plane.register_request.assert_called_once_with(new_req)  # warmup skipped
    plane.register_receivers.assert_called_once_with([handle])
    plane.request_terminal.assert_called_once_with({"r0"})
    plane.abort_requests.assert_called_once_with({"aborted"})


@pytest.mark.parametrize("output_form", ["tuple", "omni"])
def test_capture_model_unwraps_exclude_full_and_capture_mtp(output_form):
    runner = object.__new__(OmniGPUModelRunner)
    hidden = torch.ones(1, 2)

    def original_forward():
        if output_form == "tuple":
            return hidden, {"layers": {}}
        return OmniOutput(text_hidden_states=hidden, multimodal_outputs={})

    runner.model = SimpleNamespace(forward=original_forward)
    runner._model_returns_tuple = True
    runner._exclude_full_graph = True
    runner.use_aux_hidden_state_outputs = False
    piecewise = SimpleNamespace(cg_mode=CUDAGraphMode.PIECEWISE)
    full = SimpleNamespace(cg_mode=CUDAGraphMode.FULL)
    runner.cudagraph_manager = SimpleNamespace(
        _capture_descs={CUDAGraphMode.PIECEWISE: [piecewise], CUDAGraphMode.FULL: [full]},
        _candidates={(1, 0): [piecewise, full]},
    )
    runner.model_state = SimpleNamespace(capture_mtp_graphs=MagicMock())
    runner._dispatch_mtp_batch_descriptor = MagicMock(return_value="desc")

    def assert_unwrapped(_self):
        assert torch.equal(runner.model.forward(), hidden)  # forward unwrapped during capture
        return 3

    with patch.object(GPUModelRunner, "capture_model", assert_unwrapped):
        assert runner.capture_model() == 3

    assert runner.model.forward is original_forward  # restored after capture
    assert runner.cudagraph_manager._capture_descs == {CUDAGraphMode.PIECEWISE: [piecewise]}
    runner.model_state.capture_mtp_graphs.assert_called_once_with(runner._dispatch_mtp_batch_descriptor)


@pytest.mark.parametrize(
    "parallel_config,match",
    [
        (dict(pipeline_parallel_size=2, prefill_context_parallel_size=1), "pipeline parallel"),
        (dict(pipeline_parallel_size=1, prefill_context_parallel_size=2), "prefill context parallelism"),
    ],
)
def test_mrv2_rejects_parallel_modes_at_startup(parallel_config, match):
    runner = object.__new__(OmniGPUModelRunner)
    runner.vllm_config = SimpleNamespace(parallel_config=SimpleNamespace(**parallel_config))
    with pytest.raises(NotImplementedError, match=match):
        runner._validate_parallel_support()


@pytest.mark.parametrize(
    "flag,expect_omni",
    [(None, False), ("has_preprocess", True), ("has_postprocess", True), ("have_multimodal_outputs", True)],
)
def test_init_model_state_factory_dispatches_omni_only(monkeypatch, flag, expect_omni):
    upstream = MagicMock(return_value=object())
    monkeypatch.setattr("vllm_omni.worker_v2.model_states._upstream_init_model_state", upstream)
    monkeypatch.setattr(OmniModelState, "__init__", lambda *args: None)
    cfg = SimpleNamespace(model_config=SimpleNamespace(architectures=["CustomModel"]))
    model = SimpleNamespace(**({flag: True} if flag else {}))
    state = init_omni_model_state(cfg, model, None, torch.device("cpu"))

    if expect_omni:
        assert isinstance(state, OmniModelState)
        upstream.assert_not_called()
    else:
        upstream.assert_called_once()
        assert state is upstream.return_value


def test_finish_requests_notifies_model_and_cleans_only_known_slots(monkeypatch):
    runner = _make_runner()
    calls = []
    runner.model = SimpleNamespace(on_requests_finished=lambda ids: calls.append(set(ids)))
    runner.model_state = MagicMock()
    monkeypatch.setattr(GPUModelRunner, "finish_requests", lambda *args: None)
    runner.finish_requests(SimpleNamespace(finished_req_ids={"released"}, preempted_req_ids={"r1"}))
    assert calls == [{"released"}]
    assert sorted(c.args[0] for c in runner.model_state.remove_request.call_args_list) == [0]


@pytest.mark.parametrize("stage,declared", [("thinker", False), ("custom_ar", True)])
def test_capture_contract_uses_model_declaration(stage, declared):
    runner = _make_runner()
    runner.model = SimpleNamespace(model_stage=stage, _returns_tuple=declared)
    runner._configure_cudagraph_output_contract()
    assert runner._model_returns_tuple is declared
    assert runner._exclude_full_graph is declared