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

Tests are written based on the L2AdapterInterface contract defined in base.py.
Tests only use public methods and do not access private fields.
"""

# Standard
import os
import select
import time

# Third Party
import pytest
import torch

# First Party
from lmcache.v1.distributed.api import ObjectKey
from lmcache.v1.distributed.l2_adapters.config import MockL2AdapterConfig
from lmcache.v1.distributed.l2_adapters.mock_l2_adapter import MockL2Adapter
from lmcache.v1.memory_management import (
    MemoryFormat,
    MemoryObjMetadata,
    TensorMemoryObj,
)

# =============================================================================
# Test Fixtures
# =============================================================================


def create_object_key(chunk_id: int, model_name: str = "test_model") -> ObjectKey:
    """Create a test ObjectKey with the given chunk ID."""
    return ObjectKey(
        chunk_hash=ObjectKey.IntHash2Bytes(chunk_id),
        model_name=model_name,
        kv_rank=0,
    )


def create_memory_obj(size: int = 1024, fill_value: float = 1.0) -> TensorMemoryObj:
    """Create a test TensorMemoryObj with the given size."""
    raw_data = torch.empty(size, dtype=torch.float32)
    raw_data.fill_(fill_value)
    metadata = MemoryObjMetadata(
        shape=torch.Size([size]),
        dtype=torch.float32,
        address=0,
        phy_size=size * 4,
        fmt=MemoryFormat.KV_2LTD,
        ref_count=1,
    )
    return TensorMemoryObj(raw_data, metadata, parent_allocator=None)


def wait_for_event_fd(event_fd: int, timeout: float = 5.0) -> bool:
    """
    Wait for an event fd to be signaled.

    Returns True if signaled within timeout, False otherwise.
    """
    poll = select.poll()
    poll.register(event_fd, select.POLLIN)
    events = poll.poll(timeout * 1000)  # timeout in milliseconds
    if events:
        # Read and consume the event
        try:
            os.eventfd_read(event_fd)
        except BlockingIOError:
            pass
        return True
    return False


@pytest.fixture
def adapter():
    """Create a MockL2Adapter with reasonable defaults for testing."""
    config = MockL2AdapterConfig(
        max_size_gb=0.001,  # 1 MB
        mock_bandwidth_gb=10.0,  # 10 GB/s (fast for tests)
    )
    adapter = MockL2Adapter(config)
    yield adapter
    adapter.close()


@pytest.fixture
def slow_adapter():
    """Create a MockL2Adapter with slow bandwidth for timing tests."""
    config = MockL2AdapterConfig(
        max_size_gb=0.001,  # 1 MB
        mock_bandwidth_gb=0.0001,  # Very slow bandwidth
    )
    adapter = MockL2Adapter(config)
    yield adapter
    adapter.close()


# =============================================================================
# Event Fd Interface Tests
# =============================================================================


class TestEventFdInterface:
    """Test the event fd interface methods."""

    def test_get_store_event_fd_returns_valid_fd(self, adapter):
        """get_store_event_fd should return a valid file descriptor."""
        fd = adapter.get_store_event_fd()
        assert isinstance(fd, int)
        assert fd >= 0

    def test_get_lookup_and_lock_event_fd_returns_valid_fd(self, adapter):
        """get_lookup_and_lock_event_fd should return a valid file descriptor."""
        fd = adapter.get_lookup_and_lock_event_fd()
        assert isinstance(fd, int)
        assert fd >= 0

    def test_get_load_event_fd_returns_valid_fd(self, adapter):
        """get_load_event_fd should return a valid file descriptor."""
        fd = adapter.get_load_event_fd()
        assert isinstance(fd, int)
        assert fd >= 0

    def test_event_fds_are_different(self, adapter):
        """Each operation should have a distinct event fd."""
        store_fd = adapter.get_store_event_fd()
        lookup_fd = adapter.get_lookup_and_lock_event_fd()
        load_fd = adapter.get_load_event_fd()

        assert store_fd != lookup_fd
        assert store_fd != load_fd
        assert lookup_fd != load_fd


# =============================================================================
# Store Interface Tests
# =============================================================================


class TestStoreInterface:
    """Test the store operation interface."""

    def test_submit_store_task_returns_task_id(self, adapter):
        """submit_store_task should return a valid task ID."""
        key = create_object_key(1)
        obj = create_memory_obj()

        task_id = adapter.submit_store_task([key], [obj])

        assert isinstance(task_id, int)

    def test_submit_store_task_signals_event_fd(self, adapter):
        """Store completion should signal the store event fd."""
        key = create_object_key(1)
        obj = create_memory_obj()
        store_fd = adapter.get_store_event_fd()

        adapter.submit_store_task([key], [obj])

        # Wait for the event fd to be signaled
        assert wait_for_event_fd(store_fd, timeout=5.0), (
            "Store event fd was not signaled within timeout"
        )

    def test_pop_completed_store_tasks_returns_completed(self, adapter):
        """pop_completed_store_tasks should return completed tasks
        with success status.
        """
        key = create_object_key(1)
        obj = create_memory_obj()
        store_fd = adapter.get_store_event_fd()

        task_id = adapter.submit_store_task([key], [obj])

        # Wait for completion
        assert wait_for_event_fd(store_fd, timeout=5.0)

        completed = adapter.pop_completed_store_tasks()

        assert task_id in completed
        assert completed[task_id] is True  # Should be successful

    def test_pop_completed_store_tasks_clears_completed(self, adapter):
        """pop_completed_store_tasks should clear the completed tasks."""
        key = create_object_key(1)
        obj = create_memory_obj()
        store_fd = adapter.get_store_event_fd()

        adapter.submit_store_task([key], [obj])
        wait_for_event_fd(store_fd, timeout=5.0)

        # First pop should return the task
        completed1 = adapter.pop_completed_store_tasks()
        assert len(completed1) == 1

        # Second pop should return empty
        completed2 = adapter.pop_completed_store_tasks()
        assert len(completed2) == 0

    def test_submit_multiple_store_tasks(self, adapter):
        """Multiple store tasks should each get unique task IDs."""
        keys = [create_object_key(i) for i in range(3)]
        objs = [create_memory_obj(fill_value=float(i)) for i in range(3)]
        store_fd = adapter.get_store_event_fd()

        task_ids = []
        for key, obj in zip(keys, objs, strict=False):
            task_id = adapter.submit_store_task([key], [obj])
            task_ids.append(task_id)

        # All task IDs should be unique
        assert len(set(task_ids)) == 3

        # Wait for all completions
        completed = {}
        while len(completed) < 3:
            wait_for_event_fd(store_fd, timeout=5.0)
            completed.update(adapter.pop_completed_store_tasks())

        # All tasks should be completed
        for task_id in task_ids:
            assert task_id in completed
            assert completed[task_id] is True

    def test_store_batch_of_objects(self, adapter):
        """A single store task can store multiple key-object pairs."""
        keys = [create_object_key(i) for i in range(5)]
        objs = [create_memory_obj(fill_value=float(i)) for i in range(5)]
        store_fd = adapter.get_store_event_fd()

        task_id = adapter.submit_store_task(keys, objs)

        assert wait_for_event_fd(store_fd, timeout=5.0)

        completed = adapter.pop_completed_store_tasks()
        assert task_id in completed
        assert completed[task_id] is True


# =============================================================================
# Lookup and Lock Interface Tests
# =============================================================================


class TestLookupAndLockInterface:
    """Test the lookup and lock operation interface."""

    def test_submit_lookup_returns_task_id(self, adapter):
        """submit_lookup_and_lock_task should return a valid task ID."""
        key = create_object_key(1)

        task_id = adapter.submit_lookup_and_lock_task([key])

        assert isinstance(task_id, int)

    def test_lookup_signals_event_fd(self, adapter):
        """Lookup completion should signal the lookup event fd."""
        key = create_object_key(1)
        lookup_fd = adapter.get_lookup_and_lock_event_fd()

        adapter.submit_lookup_and_lock_task([key])

        assert wait_for_event_fd(lookup_fd, timeout=5.0), (
            "Lookup event fd was not signaled within timeout"
        )

    def test_lookup_nonexistent_key_returns_bitmap_with_zeros(self, adapter):
        """Looking up a non-existent key should return a bitmap with 0."""
        key = create_object_key(999)  # Never stored
        lookup_fd = adapter.get_lookup_and_lock_event_fd()

        task_id = adapter.submit_lookup_and_lock_task([key])
        wait_for_event_fd(lookup_fd, timeout=5.0)

        bitmap = adapter.query_lookup_and_lock_result(task_id)

        assert bitmap is not None
        assert bitmap.test(0) is False  # Key not found

    def test_lookup_existing_key_returns_bitmap_with_ones(self, adapter):
        """Looking up an existing key should return a bitmap with 1."""
        key = create_object_key(1)
        obj = create_memory_obj()
        store_fd = adapter.get_store_event_fd()
        lookup_fd = adapter.get_lookup_and_lock_event_fd()

        # First store the object
        adapter.submit_store_task([key], [obj])
        wait_for_event_fd(store_fd, timeout=5.0)
        adapter.pop_completed_store_tasks()

        # Now lookup
        task_id = adapter.submit_lookup_and_lock_task([key])
        wait_for_event_fd(lookup_fd, timeout=5.0)

        bitmap = adapter.query_lookup_and_lock_result(task_id)

        assert bitmap is not None
        assert bitmap.test(0) is True  # Key found

    def test_lookup_mixed_keys(self, adapter):
        """Lookup of mixed existing/non-existing keys returns correct bitmap."""
        existing_key = create_object_key(1)
        nonexistent_key = create_object_key(999)
        obj = create_memory_obj()
        store_fd = adapter.get_store_event_fd()
        lookup_fd = adapter.get_lookup_and_lock_event_fd()

        # Store only one key
        adapter.submit_store_task([existing_key], [obj])
        wait_for_event_fd(store_fd, timeout=5.0)
        adapter.pop_completed_store_tasks()

        # Lookup both keys
        task_id = adapter.submit_lookup_and_lock_task([existing_key, nonexistent_key])
        wait_for_event_fd(lookup_fd, timeout=5.0)

        bitmap = adapter.query_lookup_and_lock_result(task_id)

        assert bitmap is not None
        assert bitmap.test(0) is True  # existing_key found
        assert bitmap.test(1) is False  # nonexistent_key not found

    def test_query_lookup_result_returns_none_for_unknown_task(self, adapter):
        """Querying an unknown task ID should return None."""
        result = adapter.query_lookup_and_lock_result(99999)
        assert result is None

    def test_query_lookup_result_clears_result(self, adapter):
        """Querying lookup result should remove it (can only query once)."""
        key = create_object_key(1)
        lookup_fd = adapter.get_lookup_and_lock_event_fd()

        task_id = adapter.submit_lookup_and_lock_task([key])
        wait_for_event_fd(lookup_fd, timeout=5.0)

        # First query returns result
        result1 = adapter.query_lookup_and_lock_result(task_id)
        assert result1 is not None

        # Second query returns None
        result2 = adapter.query_lookup_and_lock_result(task_id)
        assert result2 is None


# =============================================================================
# Unlock Interface Tests
# =============================================================================


class TestUnlockInterface:
    """Test the unlock operation interface."""

    def test_submit_unlock_does_not_raise(self, adapter):
        """submit_unlock should not raise an exception."""
        key = create_object_key(1)

        # Should not raise even for non-existent key
        adapter.submit_unlock([key])

    def test_unlock_after_lock(self, adapter):
        """Unlocking after locking should work without error."""
        key = create_object_key(1)
        obj = create_memory_obj()
        store_fd = adapter.get_store_event_fd()
        lookup_fd = adapter.get_lookup_and_lock_event_fd()

        # Store
        adapter.submit_store_task([key], [obj])
        wait_for_event_fd(store_fd, timeout=5.0)
        adapter.pop_completed_store_tasks()

        # Lookup and lock
        task_id = adapter.submit_lookup_and_lock_task([key])
        wait_for_event_fd(lookup_fd, timeout=5.0)
        adapter.query_lookup_and_lock_result(task_id)

        # Unlock should not raise
        adapter.submit_unlock([key])


# =============================================================================
# Load Interface Tests
# =============================================================================


class TestLoadInterface:
    """Test the load operation interface."""

    def test_submit_load_task_returns_task_id(self, adapter):
        """submit_load_task should return a valid task ID."""
        key = create_object_key(1)
        obj = create_memory_obj()

        task_id = adapter.submit_load_task([key], [obj])

        assert isinstance(task_id, int)

    def test_load_signals_event_fd(self, adapter):
        """Load completion should signal the load event fd."""
        key = create_object_key(1)
        obj = create_memory_obj()
        load_fd = adapter.get_load_event_fd()

        adapter.submit_load_task([key], [obj])

        assert wait_for_event_fd(load_fd, timeout=5.0), (
            "Load event fd was not signaled within timeout"
        )

    def test_load_nonexistent_key_returns_bitmap_with_zeros(self, adapter):
        """Loading a non-existent key should return a bitmap with 0."""
        key = create_object_key(999)  # Never stored
        obj = create_memory_obj()
        load_fd = adapter.get_load_event_fd()

        task_id = adapter.submit_load_task([key], [obj])
        wait_for_event_fd(load_fd, timeout=5.0)

        bitmap = adapter.query_load_result(task_id)

        assert bitmap is not None
        assert bitmap.test(0) is False  # Load failed

    def test_load_existing_key_copies_data(self, adapter):
        """Loading an existing key should copy data to the provided buffer."""
        key = create_object_key(1)
        store_obj = create_memory_obj(size=100, fill_value=42.0)
        load_obj = create_memory_obj(size=100, fill_value=0.0)
        store_fd = adapter.get_store_event_fd()
        load_fd = adapter.get_load_event_fd()

        # Store
        adapter.submit_store_task([key], [store_obj])
        wait_for_event_fd(store_fd, timeout=5.0)
        adapter.pop_completed_store_tasks()

        # Load
        task_id = adapter.submit_load_task([key], [load_obj])
        wait_for_event_fd(load_fd, timeout=5.0)

        bitmap = adapter.query_load_result(task_id)

        assert bitmap is not None
        assert bitmap.test(0) is True  # Load succeeded

        # Data should be copied
        assert torch.all(load_obj.tensor == 42.0)

    def test_query_load_result_returns_none_for_unknown_task(self, adapter):
        """Querying an unknown task ID should return None."""
        result = adapter.query_load_result(99999)
        assert result is None

    def test_query_load_result_clears_result(self, adapter):
        """Querying load result should remove it (can only query once)."""
        key = create_object_key(1)
        obj = create_memory_obj()
        load_fd = adapter.get_load_event_fd()

        task_id = adapter.submit_load_task([key], [obj])
        wait_for_event_fd(load_fd, timeout=5.0)

        # First query returns result
        result1 = adapter.query_load_result(task_id)
        assert result1 is not None

        # Second query returns None
        result2 = adapter.query_load_result(task_id)
        assert result2 is None


# =============================================================================
# End-to-End Workflow Tests
# =============================================================================


class TestEndToEndWorkflow:
    """Test complete store-lookup-load workflows."""

    def test_store_lookup_load_workflow(self, adapter):
        """Test the complete workflow: store -> lookup -> load."""
        key = create_object_key(1)
        store_obj = create_memory_obj(size=256, fill_value=123.0)
        load_obj = create_memory_obj(size=256, fill_value=0.0)

        store_fd = adapter.get_store_event_fd()
        lookup_fd = adapter.get_lookup_and_lock_event_fd()
        load_fd = adapter.get_load_event_fd()

        # Step 1: Store
        store_task_id = adapter.submit_store_task([key], [store_obj])
        assert wait_for_event_fd(store_fd, timeout=5.0)
        completed = adapter.pop_completed_store_tasks()
        assert completed[store_task_id] is True

        # Step 2: Lookup and lock
        lookup_task_id = adapter.submit_lookup_and_lock_task([key])
        assert wait_for_event_fd(lookup_fd, timeout=5.0)
        lookup_bitmap = adapter.query_lookup_and_lock_result(lookup_task_id)
        assert lookup_bitmap.test(0) is True

        # Step 3: Load
        load_task_id = adapter.submit_load_task([key], [load_obj])
        assert wait_for_event_fd(load_fd, timeout=5.0)
        load_bitmap = adapter.query_load_result(load_task_id)
        assert load_bitmap.test(0) is True

        # Verify data
        assert torch.all(load_obj.tensor == 123.0)

        # Step 4: Unlock
        adapter.submit_unlock([key])

    def test_multiple_objects_workflow(self, adapter):
        """Test workflow with multiple objects."""
        num_objects = 5
        keys = [create_object_key(i) for i in range(num_objects)]
        store_objs = [
            create_memory_obj(size=64, fill_value=float(i * 10))
            for i in range(num_objects)
        ]
        load_objs = [
            create_memory_obj(size=64, fill_value=0.0) for _ in range(num_objects)
        ]

        store_fd = adapter.get_store_event_fd()
        lookup_fd = adapter.get_lookup_and_lock_event_fd()
        load_fd = adapter.get_load_event_fd()

        # Store all
        store_task_id = adapter.submit_store_task(keys, store_objs)
        assert wait_for_event_fd(store_fd, timeout=5.0)
        completed = adapter.pop_completed_store_tasks()
        assert completed[store_task_id] is True

        # Lookup all
        lookup_task_id = adapter.submit_lookup_and_lock_task(keys)
        assert wait_for_event_fd(lookup_fd, timeout=5.0)
        lookup_bitmap = adapter.query_lookup_and_lock_result(lookup_task_id)
        for i in range(num_objects):
            assert lookup_bitmap.test(i) is True

        # Load all
        load_task_id = adapter.submit_load_task(keys, load_objs)
        assert wait_for_event_fd(load_fd, timeout=5.0)
        load_bitmap = adapter.query_load_result(load_task_id)
        for i in range(num_objects):
            assert load_bitmap.test(i) is True
            assert torch.all(load_objs[i].tensor == float(i * 10))


# =============================================================================
# Close Interface Tests
# =============================================================================


class TestCloseInterface:
    """Test the close operation."""

    def test_close_does_not_raise(self):
        """close() should not raise an exception."""
        config = MockL2AdapterConfig(max_size_gb=0.001, mock_bandwidth_gb=10.0)
        adapter = MockL2Adapter(config)

        # Should not raise
        adapter.close()

    def test_close_after_operations(self):
        """close() should work after store/lookup/load operations."""
        config = MockL2AdapterConfig(max_size_gb=0.001, mock_bandwidth_gb=10.0)
        adapter = MockL2Adapter(config)

        key = create_object_key(1)
        obj = create_memory_obj()
        store_fd = adapter.get_store_event_fd()

        adapter.submit_store_task([key], [obj])
        wait_for_event_fd(store_fd, timeout=5.0)
        adapter.pop_completed_store_tasks()

        # Should not raise
        adapter.close()


# =============================================================================
# Bandwidth Simulation Tests
# =============================================================================


class TestBandwidthSimulation:
    """Test that bandwidth simulation affects completion timing."""

    def test_store_completion_is_delayed_by_bandwidth(self, slow_adapter):
        """Store completion should be delayed based on object size and bandwidth."""
        key = create_object_key(1)
        # Create a larger object to see noticeable delay
        obj = create_memory_obj(size=10000)  # ~40KB
        store_fd = slow_adapter.get_store_event_fd()

        start_time = time.time()
        slow_adapter.submit_store_task([key], [obj])

        # Should take some time due to bandwidth simulation
        assert wait_for_event_fd(store_fd, timeout=60.0)
        elapsed = time.time() - start_time

        # With very slow bandwidth, should take noticeable time
        # (exact time depends on bandwidth config)
        assert elapsed > 0.1  # At least some delay

    def test_fast_bandwidth_completes_quickly(self, adapter):
        """With fast bandwidth, completion should be quick."""
        key = create_object_key(1)
        obj = create_memory_obj(size=1000)
        store_fd = adapter.get_store_event_fd()

        start_time = time.time()
        adapter.submit_store_task([key], [obj])

        assert wait_for_event_fd(store_fd, timeout=5.0)
        elapsed = time.time() - start_time

        # With fast bandwidth, should complete quickly
        assert elapsed < 1.0
