#!/usr/bin/env python3
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""
Unit tests to verify vLLM KV events API compatibility.

These tests check that the vLLM KV events classes have the expected fields
that our Rust deserializers depend on. If vLLM changes their API, these tests
will fail early, before hitting runtime deserialization errors.

This test is the early warning for vLLM KV-event wire-format changes.

In the normal case, if this fails, update `lib/kv-router/src/zmq_wire.rs` to
match the new upstream vLLM event shape, then update this test.

That file is Dynamo's compatibility layer for vLLM KV events:
- it decodes vLLM's msgpack `array_like=True` wire format
- it handles field order changes in `BlockStored` / `BlockRemoved` / `EventBatch`
- it translates upstream `extra_keys` into Dynamo's internal `block_mm_infos`

Only touch consolidator files if we explicitly need the consolidator publisher
to preserve and republish a new upstream field.
"""

import importlib

import pytest

# Import vllm first to ensure it's properly loaded before accessing submodules.
# This works around potential issues with pytest's import machinery.
_vllm = importlib.import_module("vllm")
_kv_events = importlib.import_module("vllm.distributed.kv_events")

# Re-export the classes we need for tests
BlockStored = _kv_events.BlockStored
BlockRemoved = _kv_events.BlockRemoved
EventBatch = _kv_events.EventBatch
KVCacheEvent = _kv_events.KVCacheEvent

pytestmark = [
    pytest.mark.vllm,
    pytest.mark.unit,
    pytest.mark.gpu_0,
    pytest.mark.pre_merge,
]


def _has_group_idx(event_cls):
    return "group_idx" in event_cls.__struct_fields__


class TestVllmKvEventsApi:
    """Test vLLM KV events API compatibility."""

    def test_block_stored_fields(self):
        """Verify BlockStored has expected fields in expected order.

        The Rust deserializer expects these fields in this exact order:
        1. block_hashes
        2. parent_block_hash
        3. token_ids
        4. block_size
        5. lora_id
        6. medium
        7. lora_name (added in vLLM 0.14.0)
        8. extra_keys (added in vLLM 0.17.0)
        9. group_idx (added for hybrid KV cache groups; optional for older vLLM)

        If vLLM adds/removes/reorders fields, this test will fail.
        """
        expected_fields = [
            "block_hashes",
            "parent_block_hash",
            "token_ids",
            "block_size",
            "lora_id",
            "medium",
            "lora_name",
            "extra_keys",
        ]
        if _has_group_idx(BlockStored):
            expected_fields.append("group_idx")
        expected_fields = tuple(expected_fields)

        actual_fields = BlockStored.__struct_fields__
        assert actual_fields == expected_fields, (
            f"BlockStored fields changed!\n"
            f"Expected: {expected_fields}\n"
            f"Actual:   {actual_fields}\n"
            f"Required follow-up:\n"
            f"  - Update lib/kv-router/src/zmq_wire.rs to match the new BlockStored wire format.\n"
            f"  - Update this test's expected_fields and msgpack position checks.\n"
            f"  - If needed, add or update a regression test in lib/llm/src/kv_router/publisher.rs."
        )

    def test_block_removed_fields(self):
        """Verify BlockRemoved has expected fields in expected order."""
        expected_fields = [
            "block_hashes",
            "medium",
        ]
        if _has_group_idx(BlockRemoved):
            expected_fields.append("group_idx")
        expected_fields = tuple(expected_fields)

        actual_fields = BlockRemoved.__struct_fields__
        assert actual_fields == expected_fields, (
            f"BlockRemoved fields changed!\n"
            f"Expected: {expected_fields}\n"
            f"Actual:   {actual_fields}\n"
            f"Required follow-up:\n"
            f"  - Update lib/kv-router/src/zmq_wire.rs RawKvEvent::BlockRemoved seq deserializer.\n"
            f"  - Update this test's expected_fields."
        )

    def test_event_batch_fields(self):
        """Verify EventBatch/KVEventBatch has expected fields."""
        expected_fields = (
            "ts",
            "events",
            "data_parallel_rank",
        )

        actual_fields = EventBatch.__struct_fields__
        assert actual_fields == expected_fields, (
            f"EventBatch fields changed!\n"
            f"Expected: {expected_fields}\n"
            f"Actual:   {actual_fields}\n"
            f"Required follow-up:\n"
            f"  - Update lib/kv-router/src/zmq_wire.rs KvEventBatch Deserialize impl.\n"
            f"  - Update subscriber.rs VllmEventBatch tuple if batch field order changes.\n"
            f"  - Update this test's expected_fields."
        )

    def test_kv_cache_event_uses_array_like(self):
        """Verify KVCacheEvent uses array_like=True serialization.

        Our Rust deserializers expect msgpack arrays, not objects.
        If this changes, deserialization will break.
        """
        # msgspec structs with array_like=True have this attribute
        struct_config = getattr(KVCacheEvent, "__struct_config__", None)
        assert struct_config is not None, "KVCacheEvent is not a msgspec Struct"
        assert struct_config.array_like is True, (
            "KVCacheEvent no longer uses array_like=True! "
            "This will break Rust deserialization."
        )

    def test_kv_cache_event_uses_tag(self):
        """Verify KVCacheEvent uses tag=True for variant identification.

        The tag (e.g., 'BlockStored') is the first element in the msgpack array.
        """
        struct_config = getattr(KVCacheEvent, "__struct_config__", None)
        assert struct_config is not None, "KVCacheEvent is not a msgspec Struct"
        # When tag=True is set, struct_config.tag contains the tag string (class name)
        # or True. A falsy value (None/False) means no tagging.
        assert struct_config.tag, (
            "KVCacheEvent no longer uses tag=True! "
            "This will break Rust deserialization."
        )

    def test_block_stored_serialization_format(self):
        """Verify BlockStored serializes to expected msgpack array format.

        This is the ultimate test - if the serialized format changes,
        Rust deserialization will fail.
        """
        import msgspec

        event_kwargs = {
            "block_hashes": [123, 456],
            "parent_block_hash": 789,
            "token_ids": [1, 2, 3, 4],
            "block_size": 16,
            "lora_id": None,
            "medium": "GPU",
            "lora_name": None,
            "extra_keys": None,
        }
        if _has_group_idx(BlockStored):
            event_kwargs["group_idx"] = 0
        event = BlockStored(**event_kwargs)

        encoded = msgspec.msgpack.encode(event)
        decoded = msgspec.msgpack.decode(encoded)

        # Should be an array with tag as first element
        assert isinstance(decoded, list), f"Expected list, got {type(decoded)}"
        assert (
            decoded[0] == "BlockStored"
        ), f"Expected tag 'BlockStored', got {decoded[0]}"

        expected_len = 10 if _has_group_idx(BlockStored) else 9
        assert len(decoded) == expected_len, (
            f"Expected {expected_len} elements, got {len(decoded)}.\n"
            f"Decoded: {decoded}\n"
            f"If field count changed, update Rust deserializers."
        )

        # Verify field positions
        assert decoded[1] == [123, 456], f"block_hashes at wrong position: {decoded[1]}"
        assert decoded[2] == 789, f"parent_block_hash at wrong position: {decoded[2]}"
        assert decoded[3] == [1, 2, 3, 4], f"token_ids at wrong position: {decoded[3]}"
        assert decoded[4] == 16, f"block_size at wrong position: {decoded[4]}"
        assert decoded[5] is None, f"lora_id at wrong position: {decoded[5]}"
        assert decoded[6] == "GPU", f"medium at wrong position: {decoded[6]}"
        assert decoded[7] is None, f"lora_name at wrong position: {decoded[7]}"
        assert decoded[8] is None, f"extra_keys at wrong position: {decoded[8]}"
        if _has_group_idx(BlockStored):
            assert decoded[9] == 0, f"group_idx at wrong position: {decoded[9]}"

    def test_block_stored_tuple_extra_keys_serialization_format(self):
        """Verify multimodal tuple extra_keys keep the vLLM 0.19 wire shape."""
        import msgspec

        mm_hash = "0123456789abcdef00112233445566778899aabbccddeefffedcba9876543210"
        event_kwargs = {
            "block_hashes": [123],
            "parent_block_hash": None,
            "token_ids": [1, 2, 3, 4],
            "block_size": 16,
            "lora_id": None,
            "medium": "GPU",
            "lora_name": None,
            "extra_keys": [((mm_hash, 7),)],
        }
        if _has_group_idx(BlockStored):
            event_kwargs["group_idx"] = 0
        event = BlockStored(**event_kwargs)

        decoded = msgspec.msgpack.decode(msgspec.msgpack.encode(event))

        assert decoded[0] == "BlockStored"
        assert decoded[8] == [[[mm_hash, 7]]], (
            "vLLM multimodal extra_keys no longer serialize as nested tuple/list "
            f"payloads. Decoded: {decoded[8]!r}"
        )
        if _has_group_idx(BlockStored):
            assert decoded[9] == 0, f"group_idx at wrong position: {decoded[9]}"

    def test_block_removed_serialization_format(self):
        """Verify BlockRemoved serializes to expected msgpack array format."""
        import msgspec

        event_kwargs = {
            "block_hashes": [123, 456],
            "medium": "GPU",
        }
        if _has_group_idx(BlockRemoved):
            event_kwargs["group_idx"] = 0
        event = BlockRemoved(**event_kwargs)

        decoded = msgspec.msgpack.decode(msgspec.msgpack.encode(event))

        assert decoded[0] == "BlockRemoved"
        expected_len = 4 if _has_group_idx(BlockRemoved) else 3
        assert len(decoded) == expected_len, (
            f"Expected {expected_len} elements, got {len(decoded)}.\n"
            f"Decoded: {decoded}\n"
            f"If field count changed, update Rust deserializers."
        )
        assert decoded[1] == [123, 456], f"block_hashes at wrong position: {decoded[1]}"
        assert decoded[2] == "GPU", f"medium at wrong position: {decoded[2]}"
        if _has_group_idx(BlockRemoved):
            assert decoded[3] == 0, f"group_idx at wrong position: {decoded[3]}"
