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

import asyncio
import queue
import threading
from types import SimpleNamespace
from unittest.mock import Mock, patch

import pytest

from vllm_omni.diffusion import diffusion_engine as diffusion_engine_module
from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched import DiffusionRequestStatus, RequestScheduler
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput

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


def _make_request(request_id: str) -> OmniDiffusionRequest:
    return OmniDiffusionRequest(
        prompt=f"prompt_{request_id}",
        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
        request_id=request_id,
    )


def _make_engine() -> DiffusionEngine:
    engine = DiffusionEngine.__new__(DiffusionEngine)
    engine.od_config = SimpleNamespace(distributed_executor_backend=None)
    engine.scheduler = RequestScheduler()
    engine.scheduler.initialize(SimpleNamespace())
    engine.executor = SimpleNamespace(shutdown=Mock())
    engine._rpc_lock = threading.RLock()
    engine._cv = threading.Condition(engine._rpc_lock)
    engine._out_streams = {}
    engine._closed = False
    engine._shutting_down = False
    engine._shutdown_complete = False
    engine.abort_queue = queue.Queue()
    engine._loop_started = False
    engine.stop_event = None
    engine.worker_thread = None
    return engine


def test_close_completes_pending_output_streams() -> None:
    engine = _make_engine()
    event_loop = asyncio.new_event_loop()
    try:
        engine.main_loop = event_loop
        queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue()
        engine._out_streams["pending-stream"] = queue

        engine.close()

        output = queue.get_nowait()
        assert output.error == "DiffusionEngine is closed."
        assert output.finished is True
    finally:
        event_loop.close()


def test_emit_finished_outputs_finalizes_already_drained_waiter() -> None:
    class RacingOutQueue(dict):
        def get(self, key, default=None):
            return default

    engine = _make_engine()
    request_id = engine.scheduler.add_request(_make_request("pending-req"))
    engine.scheduler.finish_requests(request_id, DiffusionRequestStatus.FINISHED_ABORTED)
    engine._out_streams = RacingOutQueue()

    engine._emit_finished_outputs({request_id})

    assert engine.scheduler.get_request_state(request_id) is None


def test_emit_step_outputs_finalizes_finished_request_without_stream() -> None:
    engine = _make_engine()
    engine.execution_mode = DiffusionExecutionMode.STEP_BATCH
    request_id = engine.scheduler.add_request(_make_request("step-drained"))
    engine.scheduler.finish_requests(request_id, DiffusionRequestStatus.FINISHED_ABORTED)

    engine._emit_outputs({request_id}, [request_id], SimpleNamespace(get_request_output=lambda _request_id: None))

    assert engine.scheduler.get_request_state(request_id) is None


def _make_kv_cleanup_engine(mode: DiffusionKVCacheMode) -> tuple[DiffusionEngine, Mock]:
    engine = DiffusionEngine.__new__(DiffusionEngine)
    engine.od_config = SimpleNamespace(diffusion_kv_mode=mode)
    cleanup = Mock()
    engine.executor = SimpleNamespace(remove_diffusion_kv_requests=cleanup)
    return engine, cleanup


def test_paged_terminal_output_clears_worker_rows() -> None:
    engine, cleanup = _make_kv_cleanup_engine(DiffusionKVCacheMode.PAGED_SCHEDULER)
    engine._finalize_finished_request = lambda request_id, *_args: request_id
    engine._put_output = lambda *_args: None

    engine._emit_finished_outputs({"req-0", "req-1"})

    cleanup.assert_called_once()
    assert set(cleanup.call_args.args[0]) == {"req-0", "req-1"}


def test_paged_abort_clears_worker_rows_before_scheduler_free() -> None:
    engine, cleanup = _make_kv_cleanup_engine(DiffusionKVCacheMode.PAGED_SCHEDULER)
    engine.scheduler = SimpleNamespace(
        get_request_state=lambda _request_id: None,
        finish_requests=Mock(side_effect=AssertionError("unknown request must not be logically freed")),
    )

    engine._abort_requests(["req-idle", "req-idle"])

    cleanup.assert_called_once_with(["req-idle"])
    engine.scheduler.finish_requests.assert_not_called()


def test_dense_terminal_and_abort_paths_skip_worker_row_cleanup() -> None:
    engine, cleanup = _make_kv_cleanup_engine(DiffusionKVCacheMode.DENSE_LEGACY)
    engine._finalize_finished_request = lambda request_id, *_args: request_id
    engine._put_output = lambda *_args: None
    engine.scheduler = SimpleNamespace(get_request_state=lambda _request_id: None)

    engine._emit_finished_outputs({"req-0"})
    engine._abort_requests(["req-idle"])

    cleanup.assert_not_called()


def test_init_accepts_custom_scheduler(monkeypatch: pytest.MonkeyPatch) -> None:
    od_config = SimpleNamespace(
        custom_pipeline_args=None,
        model_class_name="CustomSchedulerPipeline",
        streaming_output=False,
    )
    custom_scheduler = RequestScheduler()
    fake_executor = SimpleNamespace(
        execute_request=Mock(),
        execute_batch=Mock(),
        execute_step=Mock(),
    )

    monkeypatch.setattr(
        "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
        lambda *args, **kwargs: None,
    )
    monkeypatch.setattr(
        "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
        lambda *args, **kwargs: None,
    )
    monkeypatch.setattr(
        "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
        lambda *args, **kwargs: Mock(return_value=fake_executor),
    )
    monkeypatch.setattr(
        "vllm_omni.diffusion.diffusion_engine.supports_request_batch",
        lambda *args, **kwargs: False,
    )

    engine = DiffusionEngine(od_config, scheduler=custom_scheduler)

    assert engine.scheduler is custom_scheduler


def test_scheduler_initialization_failure_closes_scheduler_and_executor(monkeypatch: pytest.MonkeyPatch) -> None:
    od_config = SimpleNamespace(
        custom_pipeline_args=None,
        model_class_name="SchedulerInitializationFailurePipeline",
        streaming_output=False,
    )
    initialization_error = RuntimeError("Scheduler initialization failed")
    custom_scheduler = SimpleNamespace(
        initialize=Mock(side_effect=initialization_error),
        close=Mock(),
    )
    fake_executor = SimpleNamespace(
        shutdown=Mock(),
    )

    monkeypatch.setattr(
        "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
        lambda *args, **kwargs: None,
    )
    monkeypatch.setattr(
        "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
        lambda *args, **kwargs: None,
    )
    monkeypatch.setattr(
        "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
        lambda *args, **kwargs: Mock(return_value=fake_executor),
    )
    monkeypatch.setattr(
        "vllm_omni.diffusion.diffusion_engine.supports_request_batch",
        lambda *args, **kwargs: False,
    )

    with pytest.raises(RuntimeError) as exc_info:
        DiffusionEngine(od_config, scheduler=custom_scheduler)

    assert exc_info.value is initialization_error
    custom_scheduler.initialize.assert_called_once_with(od_config)
    custom_scheduler.close.assert_called_once_with()


def test_init_shuts_down_executor_when_kv_control_plane_initialization_fails(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    fake_executor = SimpleNamespace(shutdown=Mock())
    monkeypatch.setattr(DiffusionEngine, "_init_process_hooks", lambda self, config: None)
    monkeypatch.setattr(
        DiffusionEngine,
        "_resolve_execution_mode",
        lambda self, config: DiffusionExecutionMode.REQUEST_BATCH,
    )
    monkeypatch.setattr(
        DiffusionEngine, "_init_executor", lambda self, config: setattr(self, "executor", fake_executor)
    )
    monkeypatch.setattr(
        "vllm_omni.diffusion.diffusion_engine.initialize_diffusion_kv_control_plane",
        Mock(side_effect=RuntimeError("KV initialization failed")),
    )

    with pytest.raises(RuntimeError, match="KV initialization failed"):
        DiffusionEngine(SimpleNamespace())

    fake_executor.shutdown.assert_called_once_with()


def test_paged_init_profiles_before_scheduler_initialization(monkeypatch: pytest.MonkeyPatch) -> None:
    events = []
    expected_profile_requests = [object()]
    fake_executor = SimpleNamespace(shutdown=Mock())
    kv_cache_config = object()
    kv_vllm_config = object()
    od_config = SimpleNamespace(diffusion_kv_mode="paged_scheduler")

    monkeypatch.setattr(DiffusionEngine, "_init_process_hooks", lambda self, config: None)
    monkeypatch.setattr(
        DiffusionEngine,
        "_resolve_execution_mode",
        lambda self, config: DiffusionExecutionMode.REQUEST_BATCH,
    )
    monkeypatch.setattr(
        DiffusionEngine,
        "_init_executor",
        lambda self, config: setattr(self, "executor", fake_executor),
    )
    monkeypatch.setattr(
        DiffusionEngine,
        "_prepare_diffusion_kv_profile_requests",
        lambda self: events.append("prepare-profile") or expected_profile_requests,
    )

    def initialize(executor, config, *, profile_requests):
        events.append("profile-workers")
        assert executor is fake_executor
        assert config is od_config
        assert profile_requests is expected_profile_requests
        return kv_cache_config, 16, 16, kv_vllm_config

    monkeypatch.setattr(
        "vllm_omni.diffusion.diffusion_engine.initialize_diffusion_kv_control_plane",
        initialize,
    )

    def init_scheduler(self, config, scheduler, *args, **kwargs):
        events.append("initialize-scheduler")
        self.scheduler = SimpleNamespace(close=Mock())

    monkeypatch.setattr(DiffusionEngine, "_init_scheduler", init_scheduler)
    monkeypatch.setattr(DiffusionEngine, "_init_runtime_state", lambda self: None)
    monkeypatch.setattr(DiffusionEngine, "_init_execute_fn", lambda self: None)
    monkeypatch.setattr(DiffusionEngine, "_log_execution_mode", lambda self, config: None)

    DiffusionEngine(od_config)

    assert events == ["prepare-profile", "profile-workers", "initialize-scheduler"]


@pytest.mark.asyncio
async def test_step_compatibility_wrapper_returns_final_batch() -> None:
    engine = _make_engine()
    first = [OmniRequestOutput.from_diffusion(request_id="req", images=[], finished=False)]
    final = [OmniRequestOutput.from_diffusion(request_id="req", images=[], finished=True)]

    async def _step_streaming(_request):
        yield first
        yield final

    engine.step_streaming = _step_streaming  # type: ignore[method-assign]
    with patch.object(diffusion_engine_module.logger, "warning_once") as warning_once:
        output = await engine.step(_make_request("req"))

    assert output is final
    warning_once.assert_called_once()


@pytest.mark.asyncio
async def test_async_wait_compatibility_wrapper_returns_final_output() -> None:
    engine = _make_engine()
    first = DiffusionOutput(output="chunk", finished=False)
    final = DiffusionOutput(output="final", finished=True)

    async def _stream_response(_request):
        yield first
        yield final

    engine.async_add_req_and_stream_response = _stream_response  # type: ignore[method-assign]
    with patch.object(diffusion_engine_module.logger, "warning_once") as warning_once:
        output = await engine.async_add_req_and_wait_for_response(_make_request("req"))

    assert output is final
    warning_once.assert_called_once()


def test_abort_request_id_aborts_scheduler_request() -> None:
    engine = _make_engine()
    request = _make_request("batch-parent")
    request_id = engine.scheduler.add_request(request)

    engine.abort("batch-parent")
    engine._process_aborts_queue()

    state = engine.scheduler.get_request_state(request_id)
    assert state is not None
    assert state.status == DiffusionRequestStatus.FINISHED_ABORTED


def test_close_rejects_late_async_requests() -> None:
    engine = _make_engine()
    event_loop = asyncio.new_event_loop()
    try:
        engine.main_loop = event_loop
        engine.close()

        with pytest.raises(RuntimeError, match="closed"):
            engine.add_request(_make_request("late-req"))
    finally:
        event_loop.close()


def test_close_resets_loop_started_for_dead_worker_thread() -> None:
    engine = _make_engine()
    engine._loop_started = True
    engine.worker_thread = SimpleNamespace(is_alive=Mock(return_value=False))

    engine.close()

    assert engine._loop_started is False


def test_close_defers_resource_shutdown_until_worker_thread_stops() -> None:
    engine = _make_engine()
    engine.scheduler.close = Mock()
    engine._loop_started = True
    worker_thread = SimpleNamespace(
        is_alive=Mock(side_effect=[True, True, False, False]),
        join=Mock(),
    )
    engine.worker_thread = worker_thread

    engine.close()

    worker_thread.join.assert_called_once_with(timeout=10)
    engine.scheduler.close.assert_not_called()
    engine.executor.shutdown.assert_not_called()
    assert engine._shutdown_complete is False
    assert engine._loop_started is True

    engine.close()

    engine.scheduler.close.assert_called_once()
    engine.executor.shutdown.assert_called_once()
    assert engine._shutdown_complete is True
    assert engine._loop_started is False


def test_fail_engine_does_not_reenter_when_executor_shutdown_fails() -> None:
    engine = _make_engine()
    shutdown_error = RuntimeError("worker shutdown failed")
    engine.executor.shutdown.side_effect = shutdown_error
    engine.scheduler.close = Mock()
    engine._fail_pending_rpcs = Mock()

    with pytest.raises(RuntimeError, match="worker shutdown failed"):
        engine._fail_engine(RuntimeError("engine failed"))

    engine._fail_engine(RuntimeError("follow-up failure"))

    engine.executor.shutdown.assert_called_once_with()
    engine.scheduler.close.assert_not_called()
    assert engine._closed is True
    assert engine._shutting_down is True
    assert engine._shutdown_complete is False


def test_finalize_aborted_request_drops_pending_async_output() -> None:
    """An aborted request with a pending async output must be drained (#6413).

    The aborted branch of _finalize_finished_request returns before the
    async_output_id branch, so without draining, the worker's late OUTPUT_READY
    is cached in the executor's _completed_outputs forever. The engine must tell
    the executor to drop that output id.
    """
    engine = _make_engine()
    engine.executor = SimpleNamespace(drop_output=Mock())
    request_id = engine.scheduler.add_request(_make_request("aborted-async"))
    engine.scheduler.finish_requests(request_id, DiffusionRequestStatus.FINISHED_ABORTED)

    runner_output = SimpleNamespace(result=None, async_output_id="aid-async-1")
    output = engine._finalize_finished_request(request_id, runner_output)

    assert output.aborted is True
    engine.executor.drop_output.assert_called_once_with("aid-async-1")


def test_finalize_aborted_request_without_async_output_skips_drop() -> None:
    """No async output pending -> nothing to drop; drop_output not called."""
    engine = _make_engine()
    engine.executor = SimpleNamespace(drop_output=Mock())
    request_id = engine.scheduler.add_request(_make_request("aborted-plain"))
    engine.scheduler.finish_requests(request_id, DiffusionRequestStatus.FINISHED_ABORTED)

    output = engine._finalize_finished_request(request_id)

    assert output.aborted is True
    engine.executor.drop_output.assert_not_called()


@pytest.mark.parametrize("scheduler_fails", [False, True])
def test_fail_engine_releases_cancellation_signals_after_workers_stop(scheduler_fails):
    from multiprocessing.shared_memory import SharedMemory

    from vllm_omni.diffusion.cancellation import RequestCancellationRegistry

    engine = _make_engine()
    engine._fail_pending_rpcs = Mock()
    registry = RequestCancellationRegistry()
    engine._request_cancellations = registry
    name = registry.create("running")

    def shutdown():
        # Cancellation is visible before shutdown, and late readers can attach.
        reader = SharedMemory(name=name)
        try:
            assert reader.buf[0] == 1
        finally:
            reader.close()

    engine.executor.shutdown.side_effect = shutdown
    engine.scheduler.close = Mock(side_effect=RuntimeError("scheduler cleanup failed") if scheduler_fails else None)
    try:
        if scheduler_fails:
            with pytest.raises(RuntimeError, match="scheduler cleanup failed"):
                engine._fail_engine(RuntimeError("engine failed"))
        else:
            engine._fail_engine(RuntimeError("engine failed"))
            engine.close()
            engine.executor.shutdown.assert_called_once_with()
        with pytest.raises(FileNotFoundError):
            SharedMemory(name=name)
        assert not registry._signals
    finally:
        registry.close()


def test_close_retains_cancellation_signal_until_stuck_worker_stops():
    from multiprocessing.shared_memory import SharedMemory

    from vllm_omni.diffusion.cancellation import RequestCancellationRegistry

    engine = _make_engine()
    registry = RequestCancellationRegistry()
    engine._request_cancellations = registry
    name = registry.create("running")
    engine.worker_thread = Mock()
    engine.worker_thread.is_alive.side_effect = [True, True, False, False]
    try:
        engine.close()
        engine.executor.shutdown.assert_not_called()
        reader = SharedMemory(name=name)
        try:
            assert reader.buf[0] == 1
        finally:
            reader.close()
        engine.close()
        engine.executor.shutdown.assert_called_once_with()
        with pytest.raises(FileNotFoundError):
            SharedMemory(name=name)
        assert not registry._signals
    finally:
        registry.close()