# SPDX-License-Identifier: Apache-2.0
"""
Unit tests for StorageManager.
"""

# Standard
import time

# Third Party
import pytest
import torch

# First Party
from lmcache.v1.distributed.api import (
    MemoryLayoutDesc,
    ObjectKey,
    PrefetchMode,
    TrimPolicy,
)
from lmcache.v1.distributed.config import (
    EvictionConfig,
    L1ManagerConfig,
    L1MemoryManagerConfig,
    StorageManagerConfig,
)
from lmcache.v1.distributed.l2_adapters.config import (
    L2AdaptersConfig,
)
from lmcache.v1.distributed.l2_adapters.mock_l2_adapter import MockL2AdapterConfig
from lmcache.v1.mp_observability.event import Event, EventType
from lmcache.v1.mp_observability.event_bus import EventBusConfig, init_event_bus

try:
    # First Party
    from lmcache.v1.distributed.storage_manager import StorageManager
except ImportError:
    # Skip tests if L1Manager cannot be imported
    pytest.skip(
        "Skipping because StorageManager cannot be imported", allow_module_level=True
    )

# Skip all tests in this module if CUDA is not available
pytestmark = pytest.mark.skipif(
    not torch.cuda.is_available(), reason="CUDA is not available"
)


def should_use_lazy_alloc() -> bool:
    """Determine if lazy allocation should be used based on CUDA availability."""
    return torch.cuda.is_available()


# =============================================================================
# Fixtures
# =============================================================================


@pytest.fixture
def basic_memory_config():
    """Create a basic L1MemoryManagerConfig for testing."""
    return L1MemoryManagerConfig(
        size_in_bytes=128 * 1024 * 1024,  # 128MB
        use_lazy=should_use_lazy_alloc(),
        init_size_in_bytes=64 * 1024 * 1024,  # 64MB
        align_bytes=0x1000,  # 4KB
    )


@pytest.fixture
def small_memory_config():
    """Create a small L1MemoryManagerConfig to test memory exhaustion."""
    return L1MemoryManagerConfig(
        size_in_bytes=64 * 1024 * 1024,  # 64MB
        use_lazy=should_use_lazy_alloc(),
        init_size_in_bytes=64 * 1024 * 1024,  # 64MB
        align_bytes=0x1000,
    )


@pytest.fixture
def basic_l1_config(basic_memory_config):
    """Create a basic L1ManagerConfig for testing."""
    return L1ManagerConfig(
        memory_config=basic_memory_config,
        write_ttl_seconds=600,
        read_ttl_seconds=300,
    )


@pytest.fixture
def small_l1_config(small_memory_config):
    """Create a small L1ManagerConfig to test memory exhaustion."""
    return L1ManagerConfig(
        memory_config=small_memory_config,
        write_ttl_seconds=600,
        read_ttl_seconds=300,
    )


@pytest.fixture
def basic_storage_manager_config(basic_l1_config):
    """Create a basic StorageManagerConfig for testing"""
    return StorageManagerConfig(
        l1_manager_config=basic_l1_config,
        eviction_config=EvictionConfig(
            eviction_policy="LRU",
        ),
    )


@pytest.fixture
def small_storage_manager_config(small_l1_config):
    """Create a small StorageManagerConfig to test memory exhaustion."""
    return StorageManagerConfig(
        l1_manager_config=small_l1_config,
        eviction_config=EvictionConfig(
            eviction_policy="LRU",
        ),
    )


@pytest.fixture
def basic_layout():
    """Create a basic MemoryLayoutDesc for testing."""
    return MemoryLayoutDesc(
        shapes=[torch.Size([100, 2, 512])],
        dtypes=[torch.bfloat16],
    )


@pytest.fixture
def large_layout():
    """Create a large MemoryLayoutDesc that will exhaust small memory.

    Each allocation is 8MB (2M elements * 4 bytes).
    """
    return MemoryLayoutDesc(
        shapes=[torch.Size([2048, 1024])],  # 2M elements * 4 bytes = 8MB
        dtypes=[torch.float32],
    )


def make_object_key(chunk_hash: int, model_name: str = "test_model", kv_rank: int = 0):
    """Helper to create ObjectKey instances."""
    hash_bytes = ObjectKey.IntHash2Bytes(chunk_hash)
    return ObjectKey(chunk_hash=hash_bytes, model_name=model_name, kv_rank=kv_rank)


def wait_for_condition(
    predicate,
    timeout: float = 5.0,
    poll_interval: float = 0.05,
) -> bool:
    """Poll until a predicate returns True or timeout."""
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        if predicate():
            return True
        time.sleep(poll_interval)
    return False


def wait_for_prefetch_status(
    sm: StorageManager,
    handle,
    timeout: float = 10.0,
    poll_interval: float = 0.05,
) -> int | None:
    """Poll query_prefetch_status until it returns a non-None value.

    Returns the contiguous prefix-hit count (``count_leading_ones``) of the
    found bitmap, matching the dense semantics these tests assert on.
    """
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        result = sm.query_prefetch_status(handle)
        if result is not None:
            return result.count_leading_ones()
        time.sleep(poll_interval)
    return None


def wait_for_sparse_found(
    sm: StorageManager,
    handle,
    timeout: float = 10.0,
    poll_interval: float = 0.05,
) -> set[int] | None:
    """Poll query_prefetch_status; return the found-key index set.

    For SPARSE prefetches the result bitmap is gap-tolerant, so callers read
    the full set via ``get_indices_list`` rather than ``count_leading_ones``.
    """
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        result = sm.query_prefetch_status(handle)
        if result is not None:
            return set(result.get_indices_list())
        time.sleep(poll_interval)
    return None


# =============================================================================
# Tests
# =============================================================================


class TestStorageManagerBasic:
    """Tests for basic functionality of StorageManager."""

    def test_basic_reserve_write(self, basic_storage_manager_config, basic_layout):
        """Test basic reserve and write functionality."""
        storage_manager = StorageManager(basic_storage_manager_config)

        object_key = make_object_key(chunk_hash=12345)

        # Reserve space for the object
        ret = storage_manager.reserve_write([object_key], basic_layout, mode="new")
        assert object_key in ret
        assert ret[object_key] is not None

        # Should not have any error
        storage_manager.finish_write([object_key])

        storage_manager.close()

    def test_reserve_write_multiple_keys(
        self, basic_storage_manager_config, basic_layout
    ):
        """Test reserve_write with multiple keys."""
        storage_manager = StorageManager(basic_storage_manager_config)

        keys = [make_object_key(i) for i in range(5)]

        ret = storage_manager.reserve_write(keys, basic_layout, mode="new")

        # All keys should be allocated
        assert len(ret) == len(keys)
        for key in keys:
            assert key in ret
            assert ret[key] is not None

        storage_manager.close()

    def test_reserve_write_oom(self, small_storage_manager_config, large_layout):
        """Test reserve_write raises L1Error on out-of-memory."""
        storage_manager = StorageManager(small_storage_manager_config)

        keys = [make_object_key(i) for i in range(20)]

        ret = storage_manager.reserve_write(keys, large_layout, mode="new")

        # At least some of the keys could be allocated
        assert len(ret) < len(keys)

        # If some keys were allocated, they should not be None
        for key, obj in ret.items():
            assert obj is not None

        storage_manager.close()

    def test_basic_prefetch(self, basic_storage_manager_config, basic_layout):
        """Test basic prefetch functionality."""
        storage_manager = StorageManager(basic_storage_manager_config)

        object_keys = [make_object_key(i) for i in range(5)]

        # Write keys into storage manager
        ret = storage_manager.reserve_write(object_keys, basic_layout, mode="new")
        for key in object_keys:
            assert key in ret
            assert ret[key] is not None
        storage_manager.finish_write(list(ret.keys()))

        # Prefetch all the objects
        handle = storage_manager.submit_prefetch_task(object_keys, basic_layout)

        hit_count = storage_manager.query_prefetch_status(handle).count_leading_ones()
        assert hit_count is not None
        assert hit_count == len(object_keys)

        storage_manager.close()

    def test_prefetch_partial_prefix_hits(
        self, basic_storage_manager_config, basic_layout
    ):
        """Test prefetch with partial hits."""
        # 5 keys: 0, 1, 3, 4 are written, 2 is missing
        storage_manager = StorageManager(basic_storage_manager_config)

        object_keys = [make_object_key(i) for i in range(5)]

        # Write only some keys into storage manager
        keys_to_write = [object_keys[0], object_keys[1], object_keys[3], object_keys[4]]
        ret = storage_manager.reserve_write(keys_to_write, basic_layout, mode="new")
        for key in keys_to_write:
            assert key in ret
            assert ret[key] is not None
        storage_manager.finish_write(list(ret.keys()))

        # Prefetch all the objects
        handle = storage_manager.submit_prefetch_task(object_keys, basic_layout)

        hit_count = storage_manager.query_prefetch_status(handle).count_leading_ones()
        assert hit_count is not None
        assert hit_count == 2  # Only 2 keys were written

        # The last 2 keys should be "writable"
        ret = storage_manager.reserve_write(
            object_keys[3:], basic_layout, mode="update"
        )
        for key in object_keys[3:]:
            assert key in ret
            assert ret[key] is not None

        storage_manager.close()

    def test_read_prefetched_basic(self, basic_storage_manager_config, basic_layout):
        """Test reading prefetched objects."""
        storage_manager = StorageManager(basic_storage_manager_config)

        object_keys = [make_object_key(i) for i in range(3)]

        # Write keys into storage manager
        ret = storage_manager.reserve_write(object_keys, basic_layout, mode="new")
        for key in object_keys:
            assert key in ret
            assert ret[key] is not None
        storage_manager.finish_write(list(ret.keys()))

        # Prefetch all the objects
        handle = storage_manager.submit_prefetch_task(object_keys, basic_layout)

        hit_count = storage_manager.query_prefetch_status(handle).count_leading_ones()
        assert hit_count is not None
        assert hit_count == len(object_keys)

        # Read the prefetched objects
        with storage_manager.read_prefetched_results(object_keys) as retrieved_objects:
            assert retrieved_objects is not None
            assert len(retrieved_objects) == len(object_keys)

        # Finish reading
        storage_manager.finish_read_prefetched(object_keys)

        # Now the objects should be writable again
        ret = storage_manager.reserve_write(object_keys, basic_layout, mode="update")
        for key in object_keys:
            assert key in ret
            assert ret[key] is not None

        storage_manager.close()

    def test_read_prefetched_not_found(
        self, basic_storage_manager_config, basic_layout
    ):
        """Test reading prefetched objects that were not found."""
        storage_manager = StorageManager(basic_storage_manager_config)

        object_keys = [make_object_key(i) for i in range(5)]

        # Write all objects into storage manager
        ret = storage_manager.reserve_write(object_keys, basic_layout, mode="new")
        for key in object_keys:
            assert key in ret
            assert ret[key] is not None
        storage_manager.finish_write(list(ret.keys()))

        # Prefetch objects except the first one
        handle = storage_manager.submit_prefetch_task(object_keys[1:], basic_layout)
        hit_count = storage_manager.query_prefetch_status(handle).count_leading_ones()
        assert hit_count is not None
        assert hit_count == len(object_keys) - 1

        # Attempt to read all the objects, should get None
        with storage_manager.read_prefetched_results(object_keys) as retrieved_objects:
            assert retrieved_objects is None

        # Remaining 4 objects should still be writable (i.e., no dangling read locks)
        ret = storage_manager.reserve_write(
            object_keys[1:], basic_layout, mode="update"
        )
        for key in object_keys[1:]:
            assert key in ret
            assert ret[key] is not None
        storage_manager.close()


# =============================================================================
# Tests for multi-reader (count / num_readers) support
# =============================================================================


class TestStorageManagerMultiReader:
    """Tests for the num_readers / count parameters.

    These parameters allow multiple workers (e.g. MLA with
    TP > 1) to each hold an independent read lock on the same
    prefetched object.
    """

    def test_prefetch_with_extra_count(
        self, basic_storage_manager_config, basic_layout
    ):
        """submit_prefetch_task(extra_count=N-1) acquires N locks."""
        sm = StorageManager(basic_storage_manager_config)
        keys = [make_object_key(i) for i in range(3)]

        # Write keys
        ret = sm.reserve_write(keys, basic_layout, mode="new")
        assert len(ret) == len(keys)
        sm.finish_write(list(ret.keys()))

        extra_count = 2  # total = 1 + 2 = 3 locks
        handle = sm.submit_prefetch_task(keys, basic_layout, extra_count=extra_count)
        hit = sm.query_prefetch_status(handle).count_leading_ones()
        assert hit == len(keys)

        # Release with matching extra_count
        sm.finish_read_prefetched(keys, extra_count=extra_count)

        # All locks released -> objects writable again
        ret = sm.reserve_write(keys, basic_layout, mode="update")
        assert len(ret) == len(keys)

        sm.close()

    def test_finish_read_prefetched_partial_extra_count(
        self, basic_storage_manager_config, basic_layout
    ):
        """Partial extra_count release leaves locks held."""
        sm = StorageManager(basic_storage_manager_config)
        keys = [make_object_key(i) for i in range(2)]

        ret = sm.reserve_write(keys, basic_layout, mode="new")
        sm.finish_write(list(ret.keys()))

        extra_count = 3  # total = 1 + 3 = 4 locks
        handle = sm.submit_prefetch_task(keys, basic_layout, extra_count=extra_count)
        hit = sm.query_prefetch_status(handle).count_leading_ones()
        assert hit == len(keys)

        # Release 2 of 4 (1 + extra_count=1)
        sm.finish_read_prefetched(keys, extra_count=1)

        # Objects should NOT be writable (2 locks remain)
        ret = sm.reserve_write(keys, basic_layout, mode="update")
        assert len(ret) == 0

        # Release remaining 2 (1 + extra_count=1)
        sm.finish_read_prefetched(keys, extra_count=1)

        # Now writable
        ret = sm.reserve_write(keys, basic_layout, mode="update")
        assert len(ret) == len(keys)

        sm.close()

    def test_prefetch_skipped_keys_released_with_extra(
        self, basic_storage_manager_config, basic_layout
    ):
        """Non-prefix L1 hits are released with correct extra.

        Keys {0,1,3,4} exist, key 2 is missing.  The prefix
        hits are {0,1}; keys {3,4} must have their N locks
        released to avoid dangling locks.
        """
        sm = StorageManager(basic_storage_manager_config)
        all_keys = [make_object_key(i) for i in range(5)]
        existing = [all_keys[i] for i in [0, 1, 3, 4]]

        ret = sm.reserve_write(existing, basic_layout, mode="new")
        sm.finish_write(list(ret.keys()))

        extra_count = 1  # total = 1 + 1 = 2 locks
        handle = sm.submit_prefetch_task(
            all_keys, basic_layout, extra_count=extra_count
        )
        hit = sm.query_prefetch_status(handle).count_leading_ones()
        # Only prefix {0,1} count as hits
        assert hit is not None
        assert hit == 2

        # Finish the prefix hits
        sm.finish_read_prefetched(all_keys[:2], extra_count=extra_count)

        # Keys {3,4} should be writable (skipped locks released)
        ret = sm.reserve_write(
            [all_keys[3], all_keys[4]],
            basic_layout,
            mode="update",
        )
        assert len(ret) == 2

        sm.close()

    def test_extra_count_default_is_zero(
        self, basic_storage_manager_config, basic_layout
    ):
        """Default extra_count=0 behaves same as before."""
        sm = StorageManager(basic_storage_manager_config)
        keys = [make_object_key(i) for i in range(3)]

        ret = sm.reserve_write(keys, basic_layout, mode="new")
        sm.finish_write(list(ret.keys()))

        handle = sm.submit_prefetch_task(keys, basic_layout)
        hit = sm.query_prefetch_status(handle).count_leading_ones()
        assert hit == len(keys)

        # Single finish is enough
        sm.finish_read_prefetched(keys)

        ret = sm.reserve_write(keys, basic_layout, mode="update")
        assert len(ret) == len(keys)

        sm.close()


# =============================================================================
# L2 Prefetch Integration Tests
# =============================================================================


@pytest.fixture
def l2_storage_manager_config(basic_l1_config):
    """Create a StorageManagerConfig with one MockL2Adapter."""
    return StorageManagerConfig(
        l1_manager_config=basic_l1_config,
        eviction_config=EvictionConfig(
            eviction_policy="LRU",
        ),
        l2_adapter_config=L2AdaptersConfig(
            adapters=[
                MockL2AdapterConfig(max_size_gb=0.01, mock_bandwidth_gb=10.0),
            ],
        ),
    )


class TestStorageManagerL2Prefetch:
    """Tests for prefetching from L2 through StorageManager."""

    def _write_keys_and_wait_for_l2(
        self,
        sm: StorageManager,
        keys: list[ObjectKey],
        layout: MemoryLayoutDesc,
    ) -> None:
        """Write keys to L1 via StorageManager and wait for L2 store."""
        ret = sm.reserve_write(keys, layout, mode="new")
        assert len(ret) == len(keys)
        sm.finish_write(list(ret.keys()))

        # Wait for StoreController to propagate all keys to L2
        adapter = sm._l2_adapters[0]
        ok = wait_for_condition(
            lambda: all(adapter.debug_has_key(k) for k in keys),  # type: ignore
            timeout=10.0,
        )
        assert ok, "Keys should be stored in L2 by StoreController"

    def test_prefetch_from_l2(self, l2_storage_manager_config, basic_layout):
        """Write to L1 → store to L2 → clear L1 → prefetch from L2."""
        sm = StorageManager(l2_storage_manager_config)
        keys = [make_object_key(i) for i in range(5)]

        self._write_keys_and_wait_for_l2(sm, keys, basic_layout)

        # Brief sleep to let StoreController release read locks
        # after L2 store completion, then clear L1
        time.sleep(0.05)
        sm.clear()
        used, _ = sm._l1_manager.get_memory_usage()
        assert used == 0, f"L1 should be empty after clear, but {used} bytes used"

        # Prefetch — L1 has 0 hits, L2 should have all 5
        handle = sm.submit_prefetch_task(keys, basic_layout)
        hit_count = wait_for_prefetch_status(sm, handle)

        assert hit_count is not None, "Prefetch should complete"
        assert hit_count == 5, f"Expected 5 hits from L2, got {hit_count}"

        # Verify keys are read-locked in L1 after prefetch
        with sm.read_prefetched_results(keys) as objs:
            assert objs is not None
            assert len(objs) == len(keys)

        sm.finish_read_prefetched(keys)
        sm.close()

    def test_prefetch_mixed_l1_l2(self, l2_storage_manager_config, basic_layout):
        """Some keys in L1, rest in L2 → combined prefix hits."""
        sm = StorageManager(l2_storage_manager_config)

        # Use distinct hash ranges for L1-only vs L2 keys
        all_keys = [make_object_key(i) for i in range(5)]
        l2_only_keys = all_keys[2:]  # keys 2, 3, 4 only in L2

        # Write all keys to L1 (StoreController will push to L2)
        self._write_keys_and_wait_for_l2(sm, all_keys, basic_layout)

        # Delete only keys 2, 3, 4 from L1 so they must come from L2
        sm._l1_manager.delete(l2_only_keys)

        # Prefetch all 5 keys: first 2 from L1, next 3 from L2
        handle = sm.submit_prefetch_task(all_keys, basic_layout)
        hit_count = wait_for_prefetch_status(sm, handle)

        assert hit_count is not None, "Prefetch should complete"
        assert hit_count == 5, (
            f"Expected 5 combined hits (2 L1 + 3 L2), got {hit_count}"
        )

        # Clean up read locks
        sm.finish_read_prefetched(all_keys)
        sm.close()

    def test_prefetch_nothing_in_l2(self, l2_storage_manager_config, basic_layout):
        """Prefetch keys not in L2 → returns 0 L2 hits."""
        sm = StorageManager(l2_storage_manager_config)

        # Don't write anything — keys exist nowhere
        keys = [make_object_key(i) for i in range(3)]

        handle = sm.submit_prefetch_task(keys, basic_layout)
        hit_count = wait_for_prefetch_status(sm, handle)

        assert hit_count is not None, "Prefetch should complete"
        assert hit_count == 0, f"Expected 0 hits, got {hit_count}"

        sm.close()

    def test_warm_skip_l2_is_noop(self, l2_storage_manager_config, basic_layout):
        """``mode=WARM`` + ``skip_l2=True`` submits no controller request.

        Regression: the WARM branch must honor ``skip_l2`` (it previously
        ignored it, issuing L2 lookup/load work and mutating L1). The returned
        :class:`PrefetchHandle` is the public signal — no request id to track
        and no L2-sourced indices; since the controller is the only path that
        loads from L2, this also means L1 is left untouched.
        """
        sm = StorageManager(l2_storage_manager_config)
        keys = [make_object_key(i) for i in range(3)]

        # Put the keys in L2 so a non-skip warm would have something to load,
        # then clear L1 so a load would be the only way they could reappear.
        self._write_keys_and_wait_for_l2(sm, keys, basic_layout)
        sm.clear()

        handle = sm.submit_prefetch_task(
            keys, basic_layout, mode=PrefetchMode.WARM, skip_l2=True
        )

        # skip_l2 honored: the controller was never asked.
        assert handle.prefetch_request_id == -1
        assert handle.l2_orig_indices == ()
        assert handle.total_requested_keys == len(keys)

        sm.close()

    def test_prefetch_l2_partial_prefix(self, l2_storage_manager_config, basic_layout):
        """L2 has keys {0,1,3,4} but not 2 → L2 returns prefix of 2."""
        sm = StorageManager(l2_storage_manager_config)

        all_keys = [make_object_key(i) for i in range(5)]
        # Write only keys 0, 1, 3, 4 (skip key 2)
        keys_to_write = [all_keys[i] for i in [0, 1, 3, 4]]
        self._write_keys_and_wait_for_l2(sm, keys_to_write, basic_layout)

        # Brief sleep to let StoreController release read locks
        # after L2 store completion, then clear L1
        time.sleep(0.05)
        sm.clear()
        used, _ = sm._l1_manager.get_memory_usage()
        assert used == 0, f"L1 should be empty after clear, but {used} bytes used"

        handle = sm.submit_prefetch_task(all_keys, basic_layout)
        hit_count = wait_for_prefetch_status(sm, handle)

        assert hit_count is not None, "Prefetch should complete"
        assert hit_count == 2, (
            f"Expected 2 prefix hits from L2 (gap at index 2), got {hit_count}"
        )

        # Only prefix keys {0, 1} should be readable
        with sm.read_prefetched_results(all_keys[:2]) as objs:
            assert objs is not None
            assert len(objs) == 2

        sm.finish_read_prefetched(all_keys[:2])
        sm.close()

    def test_prefetch_l1_prefix_plus_l2_continuation(
        self, l2_storage_manager_config, basic_layout
    ):
        """L1 has keys {0,1}, L2 has {2,3,4} → combined prefix of 5."""
        sm = StorageManager(l2_storage_manager_config)

        all_keys = [make_object_key(i) for i in range(5)]

        # Write all keys → StoreController stores all to L2
        self._write_keys_and_wait_for_l2(sm, all_keys, basic_layout)

        # Delete keys {2,3,4} from L1 only, keeping them in L2
        sm._l1_manager.delete(all_keys[2:])

        # Prefetch: L1 prefix hits = 2 (keys 0,1), L2 loads {2,3,4} → total = 5
        handle = sm.submit_prefetch_task(all_keys, basic_layout)
        hit_count = wait_for_prefetch_status(sm, handle)

        assert hit_count is not None
        assert hit_count == 5, f"Expected 5 total hits (2 L1 + 3 L2), got {hit_count}"

        sm.finish_read_prefetched(all_keys)
        sm.close()


# =============================================================================
# Tests for LM-291 failure event production
# =============================================================================


@pytest.fixture
def captured_events():
    """Enable the global event bus and capture all L1/L2 failure events.

    Replaces the process-wide ``_global_bus`` with a fresh enabled bus,
    subscribes a callback for the three failure event types, yields the
    captured-events list, then resets the bus to disabled on teardown so
    later tests don't see leftover state.
    """
    bus = init_event_bus(EventBusConfig(enabled=True, max_queue_size=10_000))
    events: list[Event] = []

    def _capture(event: Event) -> None:
        events.append(event)

    for et in (
        EventType.L1_ALLOCATION_FAILED,
        EventType.L1_READ_FAILED,
        EventType.L2_PREFETCH_FAILED,
    ):
        bus.subscribe(et, _capture)
    bus.start()
    try:
        yield events
    finally:
        bus.stop()
        init_event_bus(EventBusConfig(enabled=False))


def _events_of_type(events: list, event_type: EventType) -> list:
    return [e for e in events if e.event_type == event_type]


class TestFailureEventProduction:
    """Verifies LM-291 health-monitoring events are published at the
    right producer call sites with the expected metadata."""

    def test_reserve_write_oom_emits_l1_allocation_failed(
        self, small_storage_manager_config, large_layout, captured_events
    ):
        """OOM during user store must publish L1_ALLOCATION_FAILED with
        during=l1_store and the OOM keys."""
        sm = StorageManager(small_storage_manager_config)
        try:
            keys = [make_object_key(i) for i in range(20)]
            sm.reserve_write(keys, large_layout, mode="new")

            # Allow drain thread to deliver the event.
            assert wait_for_condition(
                lambda: len(
                    _events_of_type(captured_events, EventType.L1_ALLOCATION_FAILED)
                )
                >= 1,
                timeout=2.0,
            )

            alloc_events = _events_of_type(
                captured_events, EventType.L1_ALLOCATION_FAILED
            )
            assert len(alloc_events) == 1
            meta = alloc_events[0].metadata
            assert meta["during"] == "l1_store"
            assert len(meta["keys"]) > 0
            # All emitted keys must be from the OOM subset of the request.
            assert set(meta["keys"]).issubset(set(keys))
        finally:
            sm.close()

    def test_unsafe_read_missing_key_emits_l1_read_failed(
        self, basic_storage_manager_config, basic_layout, captured_events
    ):
        """Deleting a key between reserve_read and unsafe_read must publish
        L1_READ_FAILED with during=l1_retrieve, reason=not_found."""
        sm = StorageManager(basic_storage_manager_config)
        try:
            keys = [make_object_key(i) for i in range(3)]

            # Write + finish_write so the keys are readable.
            ret = sm.reserve_write(keys, basic_layout, mode="new")
            assert len(ret) == len(keys)
            sm.finish_write(list(ret.keys()))

            # Prefetch to acquire read locks on all keys.
            handle = sm.submit_prefetch_task(keys, basic_layout)
            assert wait_for_prefetch_status(sm, handle) == len(keys)

            # Force a mid-read race by removing the key from L1Manager's
            # internal state dict, bypassing the lock check that
            # ``L1Manager.delete`` enforces. This simulates the exact TOCTOU
            # anomaly the metric is designed to catch: reserve_read acquired
            # a lock, but the key vanished before unsafe_read.
            del sm._l1_manager._objects[keys[1]]

            # Now attempt to read — unsafe_read should find the middle key
            # missing, emitting L1_READ_FAILED(during=l1_retrieve,
            # reason=not_found).
            with sm.read_prefetched_results(keys) as objs:
                assert objs is None  # all_good=False because middle key is gone

            assert wait_for_condition(
                lambda: len(_events_of_type(captured_events, EventType.L1_READ_FAILED))
                >= 1,
                timeout=2.0,
            )

            read_events = _events_of_type(captured_events, EventType.L1_READ_FAILED)
            assert len(read_events) == 1
            meta = read_events[0].metadata
            assert meta["during"] == "l1_retrieve"
            assert meta["reason"] == "not_found"
            assert keys[1] in meta["keys"]
        finally:
            sm.close()


class TestStorageManagerSparsePrefetch:
    """SPARSE prefetch: retain a read lock on every found key, not just the
    leading contiguous prefix."""

    def test_sparse_keeps_all_found_not_just_prefix(
        self, basic_storage_manager_config, basic_layout
    ):
        """Sparse L1 prefetch retains + read-locks every found key, including
        those past a gap (unlike the contiguous-prefix default)."""
        sm = StorageManager(basic_storage_manager_config)
        all_keys = [make_object_key(i) for i in range(5)]
        # Write {0,1,3,4}; key 2 is the gap.
        existing = [all_keys[i] for i in (0, 1, 3, 4)]
        ret = sm.reserve_write(existing, basic_layout, mode="new")
        sm.finish_write(list(ret.keys()))

        handle = sm.submit_prefetch_task(
            all_keys, basic_layout, policy=TrimPolicy.SPARSE
        )
        found = wait_for_sparse_found(sm, handle, timeout=10.0)

        # Sparse: all four found indices, NOT just the prefix {0, 1}.
        assert found == {0, 1, 3, 4}

        # Every found key is read-locked (none write-reservable).
        locked = sm.reserve_write(existing, basic_layout, mode="update")
        assert len(locked) == 0

        # Releasing the full found set frees them.
        sm.finish_read_prefetched(existing)
        freed = sm.reserve_write(existing, basic_layout, mode="update")
        assert len(freed) == len(existing)

        sm.close()

    def test_sparse_from_l2_loads_all_found(
        self, l2_storage_manager_config, basic_layout
    ):
        """Sparse prefetch from L2 loads every found key (controller skips the
        prefix-only trim), not just the prefix before a gap."""
        sm = StorageManager(l2_storage_manager_config)
        all_keys = [make_object_key(i) for i in range(5)]
        # L2 has {0,1,3,4}; gap at 2.
        existing = [all_keys[i] for i in (0, 1, 3, 4)]
        wret = sm.reserve_write(existing, basic_layout, mode="new")
        sm.finish_write(list(wret.keys()))
        adapter = sm._l2_adapters[0]
        assert wait_for_condition(
            lambda: all(adapter.debug_has_key(k) for k in existing),  # type: ignore
            timeout=10.0,
        )
        time.sleep(0.05)
        sm.clear()
        used, _ = sm._l1_manager.get_memory_usage()
        assert used == 0

        handle = sm.submit_prefetch_task(
            all_keys, basic_layout, policy=TrimPolicy.SPARSE
        )
        found = wait_for_sparse_found(sm, handle, timeout=10.0)

        # Sparse from L2: all found {0,1,3,4}, NOT the contiguous prefix {0,1}.
        assert found == {0, 1, 3, 4}

        sm.finish_read_prefetched(existing)
        sm.close()