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

import asyncio
import multiprocessing as mp
import queue
import threading
import time
from types import SimpleNamespace
from unittest.mock import MagicMock, Mock

import pytest
import torch
import zmq
from vllm.v1.engine.exceptions import EngineDeadError

from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
from vllm_omni.diffusion.sched import RequestScheduler
from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput

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


# ───────────────────────────────────────────── helpers ─────────────────────


def _tagged_output(tag: str) -> DiffusionOutput:
    """Return a ``DiffusionOutput`` identifiable by its *error* field."""
    return DiffusionOutput(output=torch.tensor([0]), error=tag)


def _mock_request(tag: str):
    """Return a lightweight request object identifiable by *tag*."""
    return SimpleNamespace(
        request_id=tag,
        prompts=[f"prompt_{tag}"],
        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
    )


def _make_executor(num_gpus: int = 1):
    """Create a ``MultiprocDiffusionExecutor`` without launching workers.

    Returns ``(executor, request_queue, result_queue)``.
    """
    od_cfg = SimpleNamespace(num_gpus=num_gpus)
    monkeypatch = pytest.MonkeyPatch()
    monkeypatch.setattr(MultiprocDiffusionExecutor, "_init_executor", lambda self: None)
    executor = MultiprocDiffusionExecutor(od_cfg)
    monkeypatch.undo()

    req_q: queue.Queue = queue.Queue()
    res_q: queue.Queue = queue.Queue()

    mock_broadcast_mq = SimpleNamespace(enqueue=req_q.put)

    mock_rmq = SimpleNamespace(dequeue=lambda timeout=None: res_q.get(timeout=timeout if timeout is not None else 10))

    executor._broadcast_mq = mock_broadcast_mq
    executor._result_mq = mock_rmq
    executor._closed = False
    executor._processes = []
    executor.is_failed = False
    executor._failure_callbacks = []
    return executor, req_q, res_q


def _make_engine(num_gpus: int = 1):
    """Create a lightweight ``DiffusionEngine`` wired to mocked executor."""
    executor, req_q, res_q = _make_executor(num_gpus)
    engine = DiffusionEngine.__new__(DiffusionEngine)
    sched = RequestScheduler()
    sched.initialize(SimpleNamespace())
    engine.scheduler = sched
    engine.executor = executor
    engine._rpc_lock = threading.RLock()
    engine._cv = threading.Condition(engine._rpc_lock)
    engine._closed = False
    engine._loop_started = False
    engine._rpc_queue = queue.Queue()
    engine.abort_queue = queue.Queue()
    engine.execute_fn = executor.execute_request
    return engine, executor, req_q, res_q


def _start_worker(req_q, res_q, count=2):
    """Simulate workers: read *count* requests from *req_q* and put
    tagged ``DiffusionOutput``s on *res_q* (FIFO order).
    """

    def _run():
        for _ in range(count):
            req = req_q.get(timeout=10)
            method = req.get("method", "")
            args = req.get("args", ())
            if method in {"generate", "execute_model"} and args and hasattr(args[0], "request_id"):
                tag = f"result_for_{args[0].request_id}"
            elif args:
                tag = f"result_for_{args[0]}"
            else:
                tag = f"result_for_{method}"
            res_q.put(_tagged_output(tag))

    t = threading.Thread(target=_run, daemon=True)
    t.start()
    return t


def _inject_interleave(executor):
    """Monkey-patch ``executor._broadcast_mq.enqueue`` so that:

    * The thread named **thread_a** *blocks* after its enqueue until the
      thread named **thread_b** has finished entirely.
    * All other threads pass through unblocked.

    Returns ``(a_enqueued: Event, b_complete: Event)`` for wiring.
    """
    a_enqueued = threading.Event()
    b_complete = threading.Event()
    orig_enqueue = executor._broadcast_mq.enqueue  # points to req_q.put

    def _controlled(item):
        orig_enqueue(item)
        if threading.current_thread().name == "thread_a":
            a_enqueued.set()  # tell B: "A has enqueued"
            b_complete.wait(5)  # block A until B finishes

    executor._broadcast_mq.enqueue = _controlled
    return a_enqueued, b_complete


# ───────────────── concurrent request execution ─────────────────


class TestConcurrentRequestExecution:
    """Concurrent request execution should not swap results."""

    def test_results_are_correctly_routed(self):
        engine, executor, req_q, res_q = _make_engine()
        a_enqueued, b_complete = _inject_interleave(executor)
        wt = _start_worker(req_q, res_q, count=2)

        results: dict[str, DiffusionOutput] = {}

        def _a():
            results["A"] = engine.add_req_and_wait_for_response(_mock_request("A"))

        def _b():
            a_enqueued.wait(5)  # wait for A to enqueue
            results["B"] = engine.add_req_and_wait_for_response(_mock_request("B"))
            b_complete.set()  # release A

        ta = threading.Thread(target=_a, name="thread_a")
        tb = threading.Thread(target=_b, name="thread_b")
        ta.start()
        tb.start()
        ta.join(10)
        tb.join(10)
        wt.join(5)

        # With correct (locked) implementation both assertions hold.
        # The bug causes them to be swapped.
        assert results["A"].error == "result_for_A"
        assert results["B"].error == "result_for_B"


# ───────────────── concurrent collective RPC ─────────────────


class TestConcurrentCollectiveRpc:
    """Concurrent ``collective_rpc()`` calls should not swap results."""

    def test_results_are_correctly_routed(self):
        engine, executor, req_q, res_q = _make_engine()
        a_enqueued, b_complete = _inject_interleave(executor)
        wt = _start_worker(req_q, res_q, count=2)

        results: dict[str, object] = {}

        def _a():
            results["A"] = engine.collective_rpc(
                "ping",
                args=("call_A",),
                unique_reply_rank=0,
            )

        def _b():
            a_enqueued.wait(5)
            results["B"] = engine.collective_rpc(
                "ping",
                args=("call_B",),
                unique_reply_rank=0,
            )
            b_complete.set()

        ta = threading.Thread(target=_a, name="thread_a")
        tb = threading.Thread(target=_b, name="thread_b")
        ta.start()
        tb.start()
        ta.join(10)
        tb.join(10)
        wt.join(5)

        assert results["A"].error == "result_for_call_A"
        assert results["B"].error == "result_for_call_B"


# ──────────── concurrent request execution and collective RPC ────────────


class TestConcurrentRequestExecutionAndCollectiveRpc:
    """Request execution and ``collective_rpc()`` should not swap results."""

    def test_results_are_correctly_routed(self):
        engine, executor, req_q, res_q = _make_engine()
        a_enqueued, b_complete = _inject_interleave(executor)
        wt = _start_worker(req_q, res_q, count=2)

        results: dict[str, object] = {}

        def _a():  # request execution path
            results["A"] = engine.add_req_and_wait_for_response(_mock_request("A"))

        def _b():  # collective_rpc path
            a_enqueued.wait(5)
            results["B"] = engine.collective_rpc(
                "ping",
                args=("call_B",),
                unique_reply_rank=0,
            )
            b_complete.set()

        ta = threading.Thread(target=_a, name="thread_a")
        tb = threading.Thread(target=_b, name="thread_b")
        ta.start()
        tb.start()
        ta.join(10)
        tb.join(10)
        wt.join(5)

        assert isinstance(results["A"], DiffusionOutput)
        assert results["A"].error == "result_for_A"
        assert results["B"].error == "result_for_call_B"


# ─────────────────────── serial operation coverage ───────────────────────


class TestSerialEngineOperations:
    """Verify correct behaviour for single-threaded (serial) usage.

    These tests must pass both **before** and **after** any concurrency fix
    is applied – they guard against regressions in the basic request path.
    """

    def test_serial_add_req_returns_correct_result(self):
        engine, _, req_q, res_q = _make_engine()
        wt = _start_worker(req_q, res_q, count=1)

        result = engine.add_req_and_wait_for_response(_mock_request("X"))
        wt.join(5)

        assert isinstance(result, DiffusionOutput)
        assert result.error == "result_for_X"

    def test_serial_add_req_multiple_sequential(self):
        engine, _, req_q, res_q = _make_engine()
        wt = _start_worker(req_q, res_q, count=3)

        for tag in ("one", "two", "three"):
            out = engine.add_req_and_wait_for_response(_mock_request(tag))
            assert out.error == f"result_for_{tag}"

        wt.join(5)

    def test_serial_collective_rpc_single_rank(self):
        engine, _, req_q, res_q = _make_engine()
        wt = _start_worker(req_q, res_q, count=1)

        result = engine.collective_rpc(
            "ping",
            args=("Y",),
            unique_reply_rank=0,
        )
        wt.join(5)

        assert result.error == "result_for_Y"

    def test_serial_collective_rpc_all_ranks(self):
        """``collective_rpc`` without *unique_reply_rank* returns a single
        response from rank 0 (only rank 0 has a result_mq).
        """
        engine, _, _, res_q = _make_engine(num_gpus=2)

        # Pre-populate one result (only rank 0 replies via result_mq)
        res_q.put(_tagged_output("rank0"))

        results = engine.collective_rpc("ping", args=("multi",))

        # Only 1 response expected since only rank 0 has result_mq
        assert len(results) == 1
        assert results[0].error == "rank0"

    def test_serial_add_req_then_collective_rpc(self):
        engine, _, req_q, res_q = _make_engine()
        wt = _start_worker(req_q, res_q, count=2)

        gen_out = engine.add_req_and_wait_for_response(_mock_request("gen"))
        rpc_out = engine.collective_rpc(
            "ping",
            args=("rpc",),
            unique_reply_rank=0,
        )
        wt.join(5)

        assert gen_out.error == "result_for_gen"
        assert rpc_out.error == "result_for_rpc"

    def test_serial_add_req_error_propagation(self):
        """``add_req`` should raise when the worker reports an error."""
        engine, _, _, res_q = _make_engine()
        # Put an error response directly
        res_q.put({"status": "error", "error": "boom"})

        out = engine.add_req_and_wait_for_response(_mock_request("fail"))

        assert isinstance(out, DiffusionOutput)
        assert out.error is not None
        assert "boom" in out.error

    def test_serial_collective_rpc_error_propagation(self):
        """``collective_rpc`` should raise when the worker reports an error."""
        engine, _, _, res_q = _make_engine()
        res_q.put({"status": "error", "error": "kaboom"})

        with pytest.raises(RuntimeError, match="kaboom"):
            engine.collective_rpc("bad", unique_reply_rank=0)

    def test_collective_rpc_closed_executor_raises(self):
        engine, executor, _, _ = _make_engine()
        executor._closed = True

        with pytest.raises(RuntimeError, match="closed"):
            engine.collective_rpc("anything")


# ───────── error handling: EngineDeadError propagation through layers ─────


class TestMultiprocExecutorRaisesEngineDeadError:
    """``collective_rpc`` raises ``EngineDeadError`` when the engine is failed."""

    def test_collective_rpc_raises_when_is_failed(self):
        executor = object.__new__(MultiprocDiffusionExecutor)
        executor._closed = False
        executor._broadcast_mq = MagicMock()
        executor._result_mq = MagicMock()
        executor._result_mq.dequeue = MagicMock(side_effect=TimeoutError)
        executor.is_failed = True

        with pytest.raises(EngineDeadError):
            executor.collective_rpc(
                "generate",
                args=(MagicMock(),),
                unique_reply_rank=0,
                exec_all_ranks=True,
            )

    def test_collective_rpc_raises_mid_dequeue_when_is_failed(self):
        """Worker dies while we are polling the dequeue loop."""
        executor, _, res_q = _make_executor()

        call_count = 0
        orig_dequeue = executor._result_mq.dequeue

        def _dying_dequeue(timeout=None):
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                executor.is_failed = True
                raise TimeoutError
            return orig_dequeue(timeout=timeout)

        executor._result_mq.dequeue = _dying_dequeue

        with pytest.raises(EngineDeadError):
            executor.collective_rpc(
                "generate",
                args=(MagicMock(),),
                unique_reply_rank=0,
                exec_all_ranks=True,
            )


class TestDiffusionEngineDeadErrorPassthrough:
    """``DiffusionEngine.add_req_and_wait_for_response`` re-raises
    ``EngineDeadError`` from executor and wraps other errors."""

    def test_engine_dead_error_propagates(self):
        engine, executor, _, _ = _make_engine()
        engine.execute_fn = Mock(side_effect=EngineDeadError())

        with pytest.raises(EngineDeadError):
            engine.add_req_and_wait_for_response(_mock_request("dead"))

    def test_runtime_error_wrapped_in_output(self):
        engine, executor, _, _ = _make_engine()
        engine.execute_fn = Mock(side_effect=RuntimeError("gpu fault"))

        out = engine.add_req_and_wait_for_response(_mock_request("fault"))
        assert isinstance(out, DiffusionOutput)
        assert "gpu fault" in out.error


class TestStageDiffusionClientErrorPropagation:
    """Error surface behaviour of ``StageDiffusionClient``.

    Uses ``object.__new__`` to construct a client without spawning a real
    subprocess, then manually sets the fields needed for each test.
    """

    def _make_client(self, *, engine_dead=False, proc_alive=True):
        from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient

        client = object.__new__(StageDiffusionClient)
        client.stage_id = 0
        client.final_output = True
        client.final_output_type = "image"
        client.default_sampling_params = None
        client.custom_process_input_func = None
        client.engine_input_source = None

        client._output_queue = asyncio.Queue()
        client._rpc_results = {}
        client._pending_rpcs = set()
        client._tasks = {}
        client._shutting_down = False
        client._engine_dead = engine_dead
        client._owns_process = True
        client._proc = MagicMock(
            is_alive=MagicMock(return_value=proc_alive),
            exitcode=1,
        )
        client._request_socket = MagicMock()
        client._response_socket = MagicMock()
        client._encoder = MagicMock()
        client._decoder = MagicMock()

        return client

    @pytest.mark.asyncio
    async def test_add_request_raises_when_dead(self):
        client = self._make_client(engine_dead=True)

        with pytest.raises(EngineDeadError):
            await client.add_request_async("req-3", "test prompt", None)

    def test_check_health_raises_when_dead(self):
        client = self._make_client(engine_dead=True)

        with pytest.raises(EngineDeadError):
            client.check_health()

    def test_check_health_ok_when_alive(self):
        client = self._make_client()
        client.check_health()

    def test_get_output_raises_engine_dead_when_dead(self):
        """When ``_engine_dead`` is True and the output queue is empty,
        ``get_diffusion_output_nowait`` must raise ``EngineDeadError``."""
        client = self._make_client(engine_dead=True)
        # Simulate _drain_responses as a no-op (no ZMQ socket)
        client._response_socket.recv.side_effect = zmq.Again

        with pytest.raises(EngineDeadError):
            client.get_diffusion_output_nowait()

    def test_get_output_returns_none_when_alive_and_empty(self):
        """When the engine is alive and the queue is empty, return None."""
        client = self._make_client()
        client._response_socket.recv.side_effect = zmq.Again

        assert client.get_diffusion_output_nowait() is None

    def test_check_health_raises_when_proc_dead(self):
        """``check_health`` detects a dead subprocess via ``_proc.is_alive()``
        and raises ``EngineDeadError``, setting ``_engine_dead`` as a
        side effect."""
        client = self._make_client(proc_alive=False)

        with pytest.raises(EngineDeadError, match="not alive"):
            client.check_health()

        assert client._engine_dead is True

    def test_get_output_raises_when_proc_dead(self):
        """When the subprocess has died (non-signal exit) and the output
        queue is empty, ``get_diffusion_output_nowait`` must raise
        ``EngineDeadError`` with the exit code."""
        client = self._make_client(proc_alive=False)
        client._response_socket.recv.side_effect = zmq.Again

        with pytest.raises(EngineDeadError, match="exit code"):
            client.get_diffusion_output_nowait()

        assert client._engine_dead is True

    def test_get_output_returns_none_on_signal_death(self):
        """When the subprocess was killed by a signal (exit code > 128),
        ``get_diffusion_output_nowait`` returns ``None`` and sets
        ``_shutting_down`` instead of raising."""
        client = self._make_client(proc_alive=False)
        client._proc.exitcode = 137  # SIGKILL (128 + 9)
        client._response_socket.recv.side_effect = zmq.Again

        result = client.get_diffusion_output_nowait()

        assert result is None
        assert client._shutting_down is True
        assert client._engine_dead is True

    def test_initialize_client_requires_replica_id(self):
        from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient

        client = object.__new__(StageDiffusionClient)
        metadata = SimpleNamespace(
            stage_id=0,
            final_output=True,
            final_output_type="image",
            default_sampling_params=None,
            requires_multimodal_data=False,
            custom_process_input_func=None,
            engine_input_source=[],
        )

        with pytest.raises(AttributeError, match="replica_id"):
            client._initialize_client(
                metadata,
                "tcp://req",
                "tcp://resp",
                proc=None,
                batch_size=1,
            )


# ───────── monitor thread & death sentinel integration tests ─────────


def _poll_flag(get_flag, *, timeout=5.0, interval=0.05) -> bool:
    """Poll until ``get_flag()`` returns True or *timeout* elapses."""
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        if get_flag():
            return True
        time.sleep(interval)
    return False


def _make_short_lived_process() -> mp.Process:
    """Spawn a real subprocess that exits immediately.

    The process must be started with ``"fork"`` (or the platform default)
    so that it can use a plain ``lambda`` as its target — ``"spawn"`` would
    fail to pickle it.
    """
    ctx = mp.get_context("fork")
    p = ctx.Process(target=lambda: None, name="ShortLivedWorker-0")
    p.start()
    return p


class TestMultiprocExecutorWorkerMonitor:
    """Integration tests for ``start_worker_monitor``.

    Uses real short-lived subprocesses so that OS-level sentinel fd
    readiness is exercised end-to-end.
    """

    def test_worker_monitor_sets_is_failed_and_calls_callbacks_on_death(self):
        """When a worker process dies, the monitor thread must:
        1. Set ``is_failed = True``
        2. Call ``shutdown()`` (which sets ``_closed = True``)
        3. Invoke all registered failure callbacks
        """
        executor = object.__new__(MultiprocDiffusionExecutor)
        executor._closed = False
        executor.is_failed = False
        executor._failure_callbacks = []
        executor._broadcast_mq = None
        executor._result_mq = None
        executor.resources = None
        # Use a no-op so shutdown() doesn't crash on None resources.
        executor._finalizer = lambda: None

        proc = _make_short_lived_process()
        executor._processes = [proc]

        callback_called = threading.Event()
        executor.register_failure_callback(callback_called.set)

        executor.start_worker_monitor()

        # Wait for the process to exit and the monitor to react.
        proc.join(5)
        assert _poll_flag(lambda: executor.is_failed), "is_failed was not set"
        assert executor._closed, "shutdown() was not called"
        assert callback_called.wait(timeout=2), "failure callback was not invoked"

    def test_worker_monitor_noop_when_already_closed(self):
        """If ``_closed`` is already True when the process dies (orderly
        shutdown), the monitor must *not* set ``is_failed``."""
        executor = object.__new__(MultiprocDiffusionExecutor)
        executor._closed = True  # already shut down
        executor.is_failed = False
        executor._failure_callbacks = []
        executor._broadcast_mq = None
        executor._result_mq = None
        executor.resources = None
        executor._finalizer = lambda: None

        proc = _make_short_lived_process()
        executor._processes = [proc]

        executor.start_worker_monitor()
        proc.join(5)

        # Give the monitor thread a chance to run (it should early-return).
        time.sleep(0.3)
        assert not executor.is_failed, "is_failed should remain False on orderly shutdown"


class TestStageDiffusionClientProcMonitor:
    """Integration test for ``StageDiffusionClient._start_proc_monitor``.

    Uses a real short-lived subprocess to verify the sentinel-based
    detection pipeline.
    """

    def test_proc_monitor_sets_engine_dead_on_process_death(self):
        """When the subprocess dies, the monitor thread must set
        ``_engine_dead = True``."""
        from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient

        client = object.__new__(StageDiffusionClient)
        client.stage_id = 0
        client._shutting_down = False
        client._engine_dead = False

        proc = _make_short_lived_process()
        client._proc = proc

        client._start_proc_monitor()
        proc.join(5)

        assert _poll_flag(lambda: client._engine_dead), "_engine_dead was not set"


class TestDrainResponsesDeathSentinel:
    """Tests for death sentinel and error routing in
    ``StageDiffusionClient._drain_responses()``.
    """

    def _make_client(self):
        from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient

        client = object.__new__(StageDiffusionClient)
        client.stage_id = 0
        client._engine_dead = False
        client._shutting_down = False
        client._output_queue = asyncio.Queue()
        client._rpc_results = {}
        client._pending_rpcs = set()
        client._response_socket = MagicMock()
        client._decoder = MagicMock()
        return client

    def test_drain_responses_sets_engine_dead_on_death_sentinel(self):
        """When ``_drain_responses`` receives the ``DIFFUSION_PROC_DEAD``
        sentinel, it must set ``_engine_dead = True`` and stop draining
        (decoder is never called)."""
        client = self._make_client()

        # First recv returns the death sentinel, second would be a normal
        # message but should never be reached.
        client._response_socket.recv.side_effect = [
            StageDiffusionProc.DIFFUSION_PROC_DEAD,
            b"should-not-be-reached",
        ]

        client._drain_responses()

        assert client._engine_dead is True
        client._decoder.decode.assert_not_called()

    def test_drain_responses_routes_error_as_omni_request_output(self):
        """When ``_drain_responses`` receives a ``{"type": "error"}`` message
        with a ``request_id``, it must place an ``OmniRequestOutput`` with
        the error on ``_output_queue``."""
        client = self._make_client()

        error_msg = {
            "type": "error",
            "request_id": "req-fail",
            "error": "gpu fault",
        }
        # First recv returns the encoded error, second raises zmq.Again.
        client._response_socket.recv.side_effect = [b"encoded-error", zmq.Again]
        client._decoder.decode.return_value = error_msg

        client._drain_responses()

        assert not client._output_queue.empty()
        output = client._output_queue.get_nowait()
        assert isinstance(output, OmniRequestOutput)
        assert output.request_id == "req-fail"
        assert output.error == "gpu fault"
        assert output.finished is True