# SPDX-License-Identifier: Apache-2.0
"""
Tests for PagedSSDCacheManager and related components.

This module tests SSD-based storage for paged KV cache blocks,
enabling larger effective cache sizes than GPU memory allows.
"""

import errno
import json
import logging
import os
import queue
import shutil
import threading
import time
from pathlib import Path
from unittest.mock import patch

import pytest

from omlx.cache.paged_ssd_cache import (
    PagedSSDBlockMetadata,
    PagedSSDCacheIndex,
    PagedSSDCacheManager,
    SharedHotCacheBudget,
    _block_turboquant_bits,
    _cache_compat_signature,
    _signature_turboquant_bits,
    _extract_tensor_bytes,
    _restore_tensor_from_bytes,
    _write_safetensors_no_mx,
    parse_size,
)


def _has_mlx() -> bool:
    """Check if MLX is available."""
    try:
        import mlx.core  # noqa: F401

        return True
    except ImportError:
        return False


class TestParseSize:
    """Tests for parse_size utility function."""

    def test_parse_bytes(self):
        """Test parsing plain bytes."""
        assert parse_size("1024") == 1024
        assert parse_size("0") == 0

    def test_parse_kb(self):
        """Test parsing kilobytes."""
        assert parse_size("1KB") == 1024
        assert parse_size("10kb") == 10 * 1024
        assert parse_size("1.5KB") == int(1.5 * 1024)

    def test_parse_mb(self):
        """Test parsing megabytes."""
        assert parse_size("1MB") == 1024**2
        assert parse_size("100mb") == 100 * 1024**2

    def test_parse_gb(self):
        """Test parsing gigabytes."""
        assert parse_size("1GB") == 1024**3
        assert parse_size("16gb") == 16 * 1024**3
        assert parse_size("0.5GB") == int(0.5 * 1024**3)

    def test_parse_tb(self):
        """Test parsing terabytes."""
        assert parse_size("1TB") == 1024**4
        assert parse_size("2tb") == 2 * 1024**4

    def test_parse_with_whitespace(self):
        """Test parsing with whitespace."""
        assert parse_size("  100MB  ") == 100 * 1024**2

    def test_invalid_format(self):
        """Test invalid format raises ValueError."""
        with pytest.raises(ValueError):
            parse_size("invalid")
        with pytest.raises(ValueError):
            parse_size("MB100")


class TestPagedSSDBlockMetadata:
    """Tests for PagedSSDBlockMetadata dataclass."""

    def test_creation(self):
        """Test creating metadata."""
        metadata = PagedSSDBlockMetadata(
            block_hash=b"test_hash_bytes_1234",
            file_path=Path("/tmp/cache/test.safetensors"),
            file_size=1024,
            token_count=64,
            created_at=time.time(),
            last_access=time.time(),
            num_layers=32,
            model_name="test-model",
        )

        assert metadata.block_hash == b"test_hash_bytes_1234"
        assert metadata.file_size == 1024
        assert metadata.token_count == 64
        assert metadata.num_layers == 32
        assert metadata.model_name == "test-model"

    def test_touch(self):
        """Test touch updates last_access."""
        metadata = PagedSSDBlockMetadata(
            block_hash=b"test_hash_bytes_1234",
            file_path=Path("/tmp/test.safetensors"),
            file_size=1024,
            token_count=64,
            created_at=1000.0,
            last_access=1000.0,
            num_layers=32,
        )

        old_access = metadata.last_access
        time.sleep(0.01)
        metadata.touch()

        assert metadata.last_access > old_access

    def test_to_dict(self):
        """Test converting to dictionary."""
        now = time.time()
        metadata = PagedSSDBlockMetadata(
            block_hash=b"test_hash_bytes_1234",
            file_path=Path("/tmp/test.safetensors"),
            file_size=1024,
            token_count=64,
            created_at=now,
            last_access=now,
            num_layers=32,
            model_name="test-model",
            block_size=2048,
            cache_signature="sig",
            layer_cache_types=["KVCache", "ArraysCache"],
            layer_meta_states=[(0,), (1, 2, 3, 4)],
        )

        d = metadata.to_dict()

        assert d["block_hash"] == b"test_hash_bytes_1234".hex()
        assert d["file_path"] == "/tmp/test.safetensors"
        assert d["file_size"] == 1024
        assert d["token_count"] == 64
        assert d["num_layers"] == 32
        assert d["model_name"] == "test-model"
        assert d["block_size"] == 2048
        assert d["cache_signature"] == "sig"
        assert d["layer_cache_types"] == ["KVCache", "ArraysCache"]
        assert d["layer_meta_states"] == [[0], [1, 2, 3, 4]]

    def test_from_dict(self):
        """Test creating from dictionary."""
        d = {
            "block_hash": b"test_hash_bytes_1234".hex(),
            "file_path": "/tmp/test.safetensors",
            "file_size": 1024,
            "token_count": 64,
            "created_at": 1000.0,
            "last_access": 1000.0,
            "num_layers": 32,
            "model_name": "test-model",
            "block_size": 2048,
            "cache_signature": "sig",
            "layer_cache_types": ["KVCache", "RotatingKVCache"],
            "layer_meta_states": [[0], [1, 2, 3, 4]],
        }

        metadata = PagedSSDBlockMetadata.from_dict(d)

        assert metadata.block_hash == b"test_hash_bytes_1234"
        assert metadata.file_path == Path("/tmp/test.safetensors")
        assert metadata.file_size == 1024
        assert metadata.block_size == 2048
        assert metadata.cache_signature == "sig"
        assert metadata.layer_cache_types == ["KVCache", "RotatingKVCache"]
        assert metadata.layer_meta_states == [(0,), (1, 2, 3, 4)]

    def test_from_dict_without_optional_fields(self):
        """Test creating from dict without optional fields."""
        d = {
            "block_hash": b"test_hash".hex(),
            "file_path": "/tmp/test.safetensors",
            "file_size": 512,
            "token_count": 32,
            "created_at": 1000.0,
            "last_access": 1000.0,
            "num_layers": 16,
        }

        metadata = PagedSSDBlockMetadata.from_dict(d)

        assert metadata.model_name == ""
        assert metadata.block_size == 0
        assert metadata.cache_signature == ""
        assert metadata.layer_cache_types is None
        assert metadata.layer_meta_states is None


class TestPagedSSDCacheIndex:
    """Tests for PagedSSDCacheIndex (in-memory index)."""

    def test_empty_index(self):
        """Test empty index."""
        index = PagedSSDCacheIndex(max_size_bytes=1024**3)
        assert index.count == 0
        assert index.total_size == 0

    def test_add(self):
        """Test adding metadata."""
        index = PagedSSDCacheIndex(max_size_bytes=1024**3)

        metadata = PagedSSDBlockMetadata(
            block_hash=b"hash1_bytes_padding",
            file_path=Path("/tmp/1.safetensors"),
            file_size=1024,
            token_count=64,
            created_at=time.time(),
            last_access=time.time(),
            num_layers=32,
        )

        index.add(metadata)

        assert index.count == 1
        assert index.total_size == 1024

    def test_add_updates_existing(self):
        """Test adding with same hash updates existing entry."""
        index = PagedSSDCacheIndex(max_size_bytes=1024**3)
        block_hash = b"same_hash_bytes_pad"

        metadata1 = PagedSSDBlockMetadata(
            block_hash=block_hash,
            file_path=Path("/tmp/1.safetensors"),
            file_size=1024,
            token_count=64,
            created_at=time.time(),
            last_access=time.time(),
            num_layers=32,
        )

        metadata2 = PagedSSDBlockMetadata(
            block_hash=block_hash,
            file_path=Path("/tmp/2.safetensors"),
            file_size=2048,
            token_count=128,
            created_at=time.time(),
            last_access=time.time(),
            num_layers=32,
        )

        index.add(metadata1)
        assert index.total_size == 1024

        index.add(metadata2)
        # Should update, not add
        assert index.count == 1
        assert index.total_size == 2048

    def test_get(self):
        """Test getting metadata by hash."""
        index = PagedSSDCacheIndex(max_size_bytes=1024**3)
        block_hash = b"test_get_hash_bytes"

        metadata = PagedSSDBlockMetadata(
            block_hash=block_hash,
            file_path=Path("/tmp/test.safetensors"),
            file_size=1024,
            token_count=64,
            created_at=time.time(),
            last_access=time.time(),
            num_layers=32,
        )

        index.add(metadata)

        retrieved = index.get(block_hash)
        assert retrieved is metadata

        # Non-existent
        assert index.get(b"nonexistent_hash_by") is None

    def test_remove(self):
        """Test removing metadata."""
        index = PagedSSDCacheIndex(max_size_bytes=1024**3)
        block_hash = b"test_remove_hash_by"

        metadata = PagedSSDBlockMetadata(
            block_hash=block_hash,
            file_path=Path("/tmp/test.safetensors"),
            file_size=1024,
            token_count=64,
            created_at=time.time(),
            last_access=time.time(),
            num_layers=32,
        )

        index.add(metadata)
        assert index.count == 1

        removed = index.remove(block_hash)
        assert removed is metadata
        assert index.count == 0
        assert index.total_size == 0

    def test_remove_nonexistent(self):
        """Test removing nonexistent entry returns None."""
        index = PagedSSDCacheIndex(max_size_bytes=1024**3)
        result = index.remove(b"nonexistent_hash_by")
        assert result is None

    def test_touch(self):
        """Test touching updates LRU order."""
        index = PagedSSDCacheIndex(max_size_bytes=1024**3)

        # Add multiple entries
        for i in range(3):
            metadata = PagedSSDBlockMetadata(
                block_hash=f"hash_{i}_bytes_padding".encode()[:20],
                file_path=Path(f"/tmp/{i}.safetensors"),
                file_size=1024,
                token_count=64,
                created_at=time.time(),
                last_access=time.time(),
                num_layers=32,
            )
            index.add(metadata)
            time.sleep(0.01)  # Ensure different access times

        # Touch first entry (should move to end of LRU)
        first_hash = b"hash_0_bytes_padding"[:20]
        index.touch(first_hash)

        # Get LRU entries - first hash should not be first anymore
        lru_entries = index.get_lru_entries(3)
        lru_hashes = [e.block_hash for e in lru_entries]
        assert lru_hashes[0] != first_hash

    def test_get_lru_entries(self):
        """Test getting LRU entries."""
        index = PagedSSDCacheIndex(max_size_bytes=1024**3)

        # Add entries
        for i in range(5):
            metadata = PagedSSDBlockMetadata(
                block_hash=f"hash_{i}_bytes_padding".encode()[:20],
                file_path=Path(f"/tmp/{i}.safetensors"),
                file_size=1024,
                token_count=64,
                created_at=time.time(),
                last_access=time.time(),
                num_layers=32,
            )
            index.add(metadata)
            time.sleep(0.001)

        lru_entries = index.get_lru_entries(3)
        assert len(lru_entries) == 3

    def test_evict_until_size(self):
        """Test evicting until size limit."""
        index = PagedSSDCacheIndex(max_size_bytes=10240)

        # Add 5 entries of 1024 bytes each = 5120 total
        for i in range(5):
            metadata = PagedSSDBlockMetadata(
                block_hash=f"hash_{i}_bytes_padding".encode()[:20],
                file_path=Path(f"/tmp/{i}.safetensors"),
                file_size=1024,
                token_count=64,
                created_at=time.time(),
                last_access=time.time(),
                num_layers=32,
            )
            index.add(metadata)

        assert index.total_size == 5120

        # Evict until size is below 3000
        evicted = index.evict_until_size(3000)

        assert len(evicted) >= 2  # At least 2 entries evicted
        assert index.total_size <= 3000

    def test_contains(self):
        """Test checking if block exists."""
        index = PagedSSDCacheIndex(max_size_bytes=1024**3)
        block_hash = b"contains_test_hash1"

        assert not index.contains(block_hash)

        metadata = PagedSSDBlockMetadata(
            block_hash=block_hash,
            file_path=Path("/tmp/test.safetensors"),
            file_size=1024,
            token_count=64,
            created_at=time.time(),
            last_access=time.time(),
            num_layers=32,
        )
        index.add(metadata)

        assert index.contains(block_hash)

    def test_properties(self):
        """Test index properties."""
        max_size = 1024**3
        index = PagedSSDCacheIndex(max_size_bytes=max_size)

        assert index.max_size == max_size
        assert index.count == 0
        assert index.total_size == 0

        # Add some entries
        for i in range(3):
            metadata = PagedSSDBlockMetadata(
                block_hash=f"hash_{i}_bytes_padding".encode()[:20],
                file_path=Path(f"/tmp/{i}.safetensors"),
                file_size=1024,
                token_count=64,
                created_at=time.time(),
                last_access=time.time(),
                num_layers=32,
            )
            index.add(metadata)

        assert index.count == 3
        assert index.total_size == 3072

    def test_get_all_hashes(self):
        """Test getting all indexed hashes."""
        index = PagedSSDCacheIndex(max_size_bytes=1024**3)

        hashes = []
        for i in range(3):
            block_hash = f"hash_{i}_bytes_padding".encode()[:20]
            hashes.append(block_hash)
            metadata = PagedSSDBlockMetadata(
                block_hash=block_hash,
                file_path=Path(f"/tmp/{i}.safetensors"),
                file_size=1024,
                token_count=64,
                created_at=time.time(),
                last_access=time.time(),
                num_layers=32,
            )
            index.add(metadata)

        all_hashes = index.get_all_hashes()
        assert len(all_hashes) == 3
        for h in hashes:
            assert h in all_hashes


class TestPagedSSDCacheManager:
    """Tests for PagedSSDCacheManager."""

    def test_initialization(self, tmp_path: Path):
        """Test manager initialization."""
        cache_dir = tmp_path / "ssd_cache"

        manager = PagedSSDCacheManager(
            cache_dir=cache_dir,
            max_size_bytes=1024**3,
        )

        assert cache_dir.exists()
        # Check subdirectories created
        for char in "0123456789abcdef":
            assert (cache_dir / char).exists()

    def test_has_block(self, tmp_path: Path):
        """Test checking if block exists."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
            expected_model_name="test-model",
            expected_block_size=64,
        )

        # Non-existent block
        assert not manager.has_block(b"nonexistent_hash_by")

    def test_delete_block(self, tmp_path: Path):
        """Test deleting a block."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
        )

        # Delete non-existent
        result = manager.delete_block(b"nonexistent_hash_by")
        assert result is False

    def test_forget_block_tracks_file_as_incompatible(self, tmp_path: Path):
        """Forgetting a compatible block keeps it visible to global eviction."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
        )

        block_hash = b"\x01" * 32
        metadata = PagedSSDBlockMetadata(
            block_hash=block_hash,
            file_path=tmp_path / "ssd_cache" / "0" / "dummy.safetensors",
            file_size=1234,
            token_count=1,
            created_at=1.0,
            last_access=1.0,
            num_layers=1,
            model_name="test-model",
        )
        manager._index.add(metadata)

        assert manager.forget_block(block_hash) is True
        assert not manager._index.contains(block_hash)
        assert manager._incompatible_index.contains(block_hash)
        assert manager._tracked_ssd_size() == metadata.file_size

    def test_clear(self, tmp_path: Path):
        """Test clearing all cache."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
        )

        count = manager.clear()
        assert count == 0  # Empty cache

    def test_get_stats(self, tmp_path: Path):
        """Test getting statistics."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
        )

        stats = manager.get_stats()

        assert stats.hits == 0
        assert stats.misses == 0
        assert stats.saves == 0
        assert stats.loads == 0
        assert stats.errors == 0

    def test_get_stats_dict(self, tmp_path: Path):
        """Test getting statistics as dictionary."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
        )

        stats_dict = manager.get_stats_dict()

        assert "cache_dir" in stats_dict
        assert "max_size" in stats_dict
        assert "total_size" in stats_dict
        assert "num_files" in stats_dict
        assert "utilization" in stats_dict

    def test_cache_manager_interface(self, tmp_path: Path):
        """Test CacheManager ABC interface."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
        )

        # Test fetch (miss)
        value, hit = manager.fetch(b"nonexistent_key_byt")
        assert hit is False
        assert value is None

        # Test evict
        result = manager.evict(b"nonexistent_key_byt")
        assert result is False

        # Test size and max_size
        assert manager.size == 0
        assert manager.max_size == 1024**3

    def test_close(self, tmp_path: Path):
        """Test closing the manager."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
        )

        # Should not raise
        manager.close()

    def test_repr(self, tmp_path: Path):
        """Test string representation."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
        )

        repr_str = repr(manager)
        assert "PagedSSDCacheManager" in repr_str
        assert "ssd_cache" in repr_str

    def test_file_path_generation(self, tmp_path: Path):
        """Test file path generation uses hash-based subdirectory."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
        )

        # Test internal path generation
        block_hash = bytes.fromhex("abc123def456" + "00" * 26)  # 32 bytes
        file_path = manager._get_file_path(block_hash)

        # First hex char of hash determines subdirectory
        assert file_path.parent.name == "a"
        assert file_path.suffix == ".safetensors"

    def test_enforce_size_limit(self, tmp_path: Path):
        """Test enforcing size limit."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
        )

        # Should return 0 when under limit
        freed = manager.enforce_size_limit()
        assert freed == 0


class TestPagedSSDCacheManagerWithMLX:
    """Tests for PagedSSDCacheManager that require MLX.

    These tests are skipped if MLX is not available.
    """

    @pytest.fixture
    def mock_mlx(self):
        """Mock MLX module for testing save/load without actual tensors."""
        try:
            import mlx.core as mx

            return mx
        except ImportError:
            pytest.skip("MLX not available")

    def test_save_and_load_block(self, tmp_path: Path, mock_mlx):
        """Test saving and loading a block with actual tensors."""
        mx = mock_mlx

        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
        )

        # Create test cache data
        block_hash = b"test_save_load_hash1"
        cache_data = [
            (mx.zeros((1, 8, 64, 64)), mx.zeros((1, 8, 64, 64)))
            for _ in range(4)  # 4 layers
        ]

        # Save
        result = manager.save_block(
            block_hash=block_hash,
            cache_data=cache_data,
            token_count=64,
            model_name="test-model",
            layer_cache_types=["KVCache"] * 4,
        )
        assert result is True
        assert manager.has_block(block_hash)

        # Load
        loaded = manager.load_block(block_hash)
        assert loaded is not None
        assert len(loaded) == 4

        # Verify shapes
        for keys, values in loaded:
            assert keys.shape == (1, 8, 64, 64)
            assert values.shape == (1, 8, 64, 64)

    def test_load_block_with_metadata(self, tmp_path: Path, mock_mlx):
        """Test loading block with metadata."""
        mx = mock_mlx

        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
        )

        block_hash = b"test_load_meta_hash"
        cache_data = [
            (mx.zeros((1, 8, 64, 64)), mx.zeros((1, 8, 64, 64))) for _ in range(2)
        ]

        manager.save_block(
            block_hash=block_hash,
            cache_data=cache_data,
            token_count=64,
            model_name="test-model",
            layer_cache_types=["KVCache", "RotatingKVCache"],
            layer_meta_states=[(0,), (1, 256, 64, 0)],
        )

        # Load with metadata
        loaded_data, loaded_meta = manager.load_block_with_metadata(block_hash)

        assert loaded_data is not None
        assert loaded_meta is not None
        assert loaded_meta["num_layers"] == 2
        assert loaded_meta["token_count"] == 64
        assert loaded_meta["model_name"] == "test-model"
        assert loaded_meta["block_size"] == 64
        assert loaded_meta["cache_signature"]
        assert loaded_meta["layer_cache_types"] == ["KVCache", "RotatingKVCache"]

    def test_save_canonicalizes_buffered_rotating_metadata(
        self, tmp_path: Path, mock_mlx
    ):
        """Transient MTP rotating wrappers should not persist in new metadata."""
        mx = mock_mlx

        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
        )

        block_hash = b"buffered_rotating_hash"
        cache_data = [
            (mx.zeros((1, 8, 64, 64)), mx.zeros((1, 8, 64, 64))),
            (mx.zeros((1, 8, 64, 64)), mx.zeros((1, 8, 64, 64))),
        ]

        manager.save_block(
            block_hash=block_hash,
            cache_data=cache_data,
            token_count=64,
            model_name="test-model",
            layer_cache_types=["KVCache", "BufferedRotatingKVCache"],
            layer_meta_states=[(64,), ("0", "1024", "2048", "64", "1984", "64")],
        )

        _, loaded_meta = manager.load_block_with_metadata(block_hash)

        assert loaded_meta is not None
        assert loaded_meta["layer_cache_types"] == ["KVCache", "RotatingKVCache"]
        signature = json.loads(loaded_meta["cache_signature"])
        assert signature["layer_cache_types"] == ["KVCache", "RotatingKVCache"]

    def test_get_block_metadata(self, tmp_path: Path, mock_mlx):
        """Test getting block metadata without loading data."""
        mx = mock_mlx

        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
        )

        block_hash = b"test_get_metadata_h"
        cache_data = [(mx.zeros((1, 8, 32, 64)), mx.zeros((1, 8, 32, 64)))]

        manager.save_block(
            block_hash=block_hash,
            cache_data=cache_data,
            token_count=32,
            model_name="test-model",
        )

        metadata = manager.get_block_metadata(block_hash)

        assert metadata is not None
        assert metadata.block_hash == block_hash
        assert metadata.token_count == 32
        assert metadata.num_layers == 1
        assert metadata.model_name == "test-model"

    def test_save_existing_block_touches(self, tmp_path: Path, mock_mlx):
        """Test saving existing block just touches (updates LRU)."""
        mx = mock_mlx

        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
        )

        block_hash = b"test_touch_existing"
        cache_data = [(mx.zeros((1, 8, 32, 64)), mx.zeros((1, 8, 32, 64)))]

        # First save
        manager.save_block(block_hash, cache_data, 32)
        initial_saves = manager._stats["saves"]

        # Second save (should just touch)
        manager.save_block(block_hash, cache_data, 32)

        # saves count should not increase (just hit)
        assert manager._stats["saves"] == initial_saves
        assert manager._stats["hits"] >= 1

    def test_save_writes_format_version(self, tmp_path: Path, mock_mlx):
        """Saved blocks tag the file with the current format version."""
        import time as time_mod

        from omlx.cache.paged_ssd_cache import _CACHE_FORMAT_VERSION

        mx = mock_mlx

        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
        )

        block_hash = b"test_format_version_save"
        cache_data = [(mx.zeros((1, 8, 32, 64)), mx.zeros((1, 8, 32, 64)))]
        assert manager.save_block(block_hash, cache_data, 32) is True

        # Wait for the background writer to flush the file to disk.
        file_path = manager._get_file_path(block_hash)
        for _ in range(50):
            if file_path.exists():
                break
            time_mod.sleep(0.1)
        assert file_path.exists(), "background writer never produced the file"

        _, file_metadata = mx.load(str(file_path), return_metadata=True)
        assert file_metadata.get("omlx_cache_format_version") == _CACHE_FORMAT_VERSION

    def test_unversioned_block_is_rejected_at_index_scan(
        self, tmp_path: Path, mock_mlx
    ):
        """Pre-fix blocks (no version marker) are skipped during scan.

        Older builds saved RotatingKVCache layers zero-padded to max_size.
        Loading those after the fix would leak zero positions into
        attention via BatchRotatingKVCache.merge(). Treat them as a cache
        miss by rejecting blocks without the format version.
        """
        mx = mock_mlx

        cache_dir = tmp_path / "ssd_cache"
        cache_dir.mkdir(parents=True, exist_ok=True)

        # Hand-write a cache file without the version tag, mirroring what
        # an old-format save_block() would produce.
        block_hash = b"\x01" * 32
        block_hash_hex = block_hash.hex()
        # Match the manager's per-prefix subdirectory layout.
        sub_dir = cache_dir / block_hash_hex[:2]
        sub_dir.mkdir(parents=True, exist_ok=True)
        legacy_file = sub_dir / f"{block_hash_hex}.safetensors"

        mx.save_safetensors(
            str(legacy_file),
            {
                "layer_0_keys": mx.zeros((1, 8, 32, 64)),
                "layer_0_values": mx.zeros((1, 8, 32, 64)),
            },
            metadata={
                # Intentionally missing omlx_cache_format_version.
                "block_hash": block_hash_hex,
                "token_count": "32",
                "num_layers": "1",
                "model_name": "legacy-model",
                "created_at": "0",
            },
        )

        manager_after_scan = PagedSSDCacheManager(
            cache_dir=cache_dir,
            max_size_bytes=1024**3,
        )

        # Index scan ran in __init__. The legacy file should not appear.
        assert not manager_after_scan.has_block(block_hash)

    def _write_versioned_fixture_block(
        self,
        cache_dir: Path,
        mx,
        block_hash: bytes,
        *,
        num_layers: int,
        model_name: str,
        block_size: int = 256,
        layer_cache_types: list[str] | None = None,
    ) -> Path:
        """Drop a minimally-valid versioned block on disk so we can exercise
        the startup scan without relying on the background writer."""
        from omlx.cache.paged_ssd_cache import _CACHE_FORMAT_VERSION

        cache_dir.mkdir(parents=True, exist_ok=True)
        block_hash_hex = block_hash.hex()
        sub_dir = cache_dir / block_hash_hex[0]
        sub_dir.mkdir(parents=True, exist_ok=True)
        file_path = sub_dir / f"{block_hash_hex}.safetensors"

        tensors = {}
        for i in range(num_layers):
            tensors[f"layer_{i}_keys"] = mx.zeros((1, 8, 32, 64))
            tensors[f"layer_{i}_values"] = mx.zeros((1, 8, 32, 64))

        if layer_cache_types is None:
            layer_cache_types = ["KVCache"] * num_layers

        mx.save_safetensors(
            str(file_path),
            tensors,
            metadata={
                "omlx_cache_format_version": _CACHE_FORMAT_VERSION,
                "block_hash": block_hash_hex,
                "token_count": "32",
                "num_layers": str(num_layers),
                "model_name": model_name,
                "block_size": str(block_size),
                "cache_signature": _cache_compat_signature(
                    model_name=model_name,
                    num_layers=num_layers,
                    block_size=block_size,
                    layer_cache_types=layer_cache_types,
                ),
                "layer_cache_types": json.dumps(layer_cache_types),
                "created_at": "0",
            },
        )
        return file_path

    def test_scan_skips_layer_count_mismatch_without_unlinking(
        self, tmp_path: Path, mock_mlx
    ):
        """Blocks with num_layers != expected_num_layers are not indexed.

        Models that change their effective layer count across versions (e.g.,
        #1404 attaching MTPModule changed 30 -> 40) should not hit the
        layer-mismatch reject path on every prefix lookup. The file is still
        left on disk so shared cache directories are non-destructive.
        """
        mx = mock_mlx
        cache_dir = tmp_path / "ssd_cache"

        stale_hash = b"\x10" + b"\x00" * 31
        fresh_hash = b"\x20" + b"\x00" * 31
        stale_path = self._write_versioned_fixture_block(
            cache_dir, mx, stale_hash, num_layers=30, model_name="qwen3.6"
        )
        fresh_path = self._write_versioned_fixture_block(
            cache_dir, mx, fresh_hash, num_layers=40, model_name="qwen3.6"
        )

        manager = PagedSSDCacheManager(
            cache_dir=cache_dir,
            max_size_bytes=1024**3,
            expected_model_name="qwen3.6",
            expected_num_layers=40,
            expected_block_size=256,
        )

        assert stale_path.exists()
        assert fresh_path.exists()
        assert not manager.has_block(stale_hash)
        assert manager.has_block(fresh_hash)

    def test_scan_skips_model_name_mismatch_without_unlinking(
        self, tmp_path: Path, mock_mlx
    ):
        """Blocks from a different model stay on disk but are not indexed."""
        mx = mock_mlx
        cache_dir = tmp_path / "ssd_cache"

        other_hash = b"\x30" + b"\x00" * 31
        match_hash = b"\x40" + b"\x00" * 31
        other_path = self._write_versioned_fixture_block(
            cache_dir, mx, other_hash, num_layers=40, model_name="llama"
        )
        match_path = self._write_versioned_fixture_block(
            cache_dir, mx, match_hash, num_layers=40, model_name="qwen3.6"
        )

        manager = PagedSSDCacheManager(
            cache_dir=cache_dir,
            max_size_bytes=1024**3,
            expected_model_name="qwen3.6",
            expected_num_layers=40,
            expected_block_size=256,
        )

        assert other_path.exists()
        assert match_path.exists()
        assert not manager.has_block(other_hash)
        assert manager.has_block(match_hash)

    def test_scan_skips_block_size_mismatch_without_unlinking(
        self, tmp_path: Path, mock_mlx
    ):
        """Blocks with another paged cache block size are not indexed."""
        mx = mock_mlx
        cache_dir = tmp_path / "ssd_cache"

        wrong_hash = b"\x41" + b"\x00" * 31
        match_hash = b"\x42" + b"\x00" * 31
        wrong_path = self._write_versioned_fixture_block(
            cache_dir,
            mx,
            wrong_hash,
            num_layers=40,
            model_name="qwen3.6",
            block_size=2048,
        )
        match_path = self._write_versioned_fixture_block(
            cache_dir,
            mx,
            match_hash,
            num_layers=40,
            model_name="qwen3.6",
            block_size=256,
        )

        manager = PagedSSDCacheManager(
            cache_dir=cache_dir,
            max_size_bytes=1024**3,
            expected_model_name="qwen3.6",
            expected_num_layers=40,
            expected_block_size=256,
        )

        assert wrong_path.exists()
        assert match_path.exists()
        assert not manager.has_block(wrong_hash)
        assert manager.has_block(match_hash)

    def test_scan_keeps_blocks_when_expected_fields_unset(
        self, tmp_path: Path, mock_mlx
    ):
        """Backwards compatibility: callers that omit the new init args see
        no behavior change. All blocks survive scan regardless of metadata."""
        mx = mock_mlx
        cache_dir = tmp_path / "ssd_cache"

        h1 = b"\x50" + b"\x00" * 31
        h2 = b"\x60" + b"\x00" * 31
        p1 = self._write_versioned_fixture_block(
            cache_dir, mx, h1, num_layers=30, model_name="a"
        )
        p2 = self._write_versioned_fixture_block(
            cache_dir, mx, h2, num_layers=40, model_name="b"
        )

        manager = PagedSSDCacheManager(
            cache_dir=cache_dir,
            max_size_bytes=1024**3,
        )

        assert p1.exists()
        assert p2.exists()
        assert manager.has_block(h1)
        assert manager.has_block(h2)

    def test_scan_logs_skipped_incompatible_count(
        self, tmp_path: Path, mock_mlx, caplog
    ):
        """The completion log line surfaces incompatible blocks skipped at scan."""
        import logging

        mx = mock_mlx
        cache_dir = tmp_path / "ssd_cache"

        for i in range(3):
            self._write_versioned_fixture_block(
                cache_dir,
                mx,
                bytes([0x70 + i]) + b"\x00" * 31,
                num_layers=30,
                model_name="old",
            )

        with caplog.at_level(logging.INFO, logger="omlx.cache.paged_ssd_cache"):
            PagedSSDCacheManager(
                cache_dir=cache_dir,
                max_size_bytes=1024**3,
                expected_model_name="old",
                expected_num_layers=40,
                expected_block_size=256,
            )

        scan_lines = [
            r.message for r in caplog.records if "SSD cache scan complete" in r.message
        ]
        assert scan_lines, "scan completion log not emitted"
        assert "skipped_incompatible=3 blocks" in scan_lines[-1]

    def test_model_switch_enforces_shared_ssd_limit_on_new_save(
        self, tmp_path: Path, mock_mlx
    ):
        """Incompatible old-model blocks count against the shared SSD budget."""
        mx = mock_mlx
        cache_dir = tmp_path / "ssd_cache"

        old_hash = b"\x80" + b"\x00" * 31
        old_path = self._write_versioned_fixture_block(
            cache_dir,
            mx,
            old_hash,
            num_layers=1,
            model_name="old-model",
        )
        old_size = old_path.stat().st_size
        max_size = old_size + old_size // 2

        manager = PagedSSDCacheManager(
            cache_dir=cache_dir,
            max_size_bytes=max_size,
            expected_model_name="new-model",
            expected_num_layers=1,
            expected_block_size=256,
        )
        try:
            assert not manager.has_block(old_hash)
            assert old_path.exists()
            assert manager._incompatible_index.total_size == old_size

            new_hash = b"\x81" + b"\x00" * 31
            cache_data = [(mx.zeros((1, 8, 32, 64)), mx.zeros((1, 8, 32, 64)))]

            assert manager.save_block(
                block_hash=new_hash,
                cache_data=cache_data,
                token_count=32,
                model_name="new-model",
                layer_cache_types=["KVCache"],
            )

            assert manager.has_block(new_hash)
            assert not old_path.exists()
            assert manager._incompatible_index.total_size == 0
            assert manager._tracked_ssd_size() <= max_size
        finally:
            manager.close()

    def test_scan_cleans_oldest_incompatible_blocks_when_over_limit(
        self, tmp_path: Path, mock_mlx
    ):
        """Startup cleanup converges an already-over-limit shared cache dir."""
        mx = mock_mlx
        cache_dir = tmp_path / "ssd_cache"

        paths = []
        for i in range(3):
            path = self._write_versioned_fixture_block(
                cache_dir,
                mx,
                bytes([0x90 + i]) + b"\x00" * 31,
                num_layers=1,
                model_name="old-model",
            )
            os.utime(path, (100 + i, 100 + i))
            paths.append(path)

        sizes = [path.stat().st_size for path in paths]
        max_size = sum(sizes) - sizes[0] + 1
        manager = PagedSSDCacheManager(
            cache_dir=cache_dir,
            max_size_bytes=max_size,
            expected_model_name="new-model",
            expected_num_layers=1,
            expected_block_size=256,
        )
        try:
            assert not paths[0].exists()
            assert paths[1].exists()
            assert paths[2].exists()
            assert manager._tracked_ssd_size() <= max_size
        finally:
            manager.close()

    def test_clear_removes_incompatible_scanned_files(self, tmp_path: Path, mock_mlx):
        """Global clear deletes compatible and incompatible tracked files."""
        mx = mock_mlx
        cache_dir = tmp_path / "ssd_cache"

        other_hash = b"\xa0" + b"\x00" * 31
        match_hash = b"\xa1" + b"\x00" * 31
        other_path = self._write_versioned_fixture_block(
            cache_dir,
            mx,
            other_hash,
            num_layers=1,
            model_name="old-model",
        )
        match_path = self._write_versioned_fixture_block(
            cache_dir,
            mx,
            match_hash,
            num_layers=1,
            model_name="new-model",
        )

        manager = PagedSSDCacheManager(
            cache_dir=cache_dir,
            max_size_bytes=1024**3,
            expected_model_name="new-model",
            expected_num_layers=1,
            expected_block_size=256,
        )
        try:
            assert manager.clear() == 2
            assert not other_path.exists()
            assert not match_path.exists()
            assert manager._tracked_ssd_count() == 0
        finally:
            manager.close()


class TestPagedSSDCacheManagerCacheList:
    """Tests for CacheList support in PagedSSDCacheManager."""

    @pytest.fixture
    def mx(self):
        """Import MLX or skip."""
        try:
            import mlx.core as mx

            return mx
        except ImportError:
            pytest.skip("MLX not available")

    @pytest.fixture
    def ssd_cache(self, tmp_path):
        """Create a PagedSSDCacheManager for testing."""
        return PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=100 * 1024**2,
        )

    def test_save_load_cache_list_block(self, ssd_cache, mx):
        """Test saving and loading a block with CacheList data."""
        block_hash = b"cache_list_test_hash"
        # Build cache_data with CacheList marker
        sub_keys1 = mx.zeros((1, 8, 32, 64))
        sub_values1 = mx.ones((1, 8, 32, 64))
        sub_keys2 = mx.zeros((1, 4, 32, 64))
        sub_values2 = mx.ones((1, 4, 32, 64))

        cache_data = [
            ("__cache_list__", [(sub_keys1, sub_values1), (sub_keys2, sub_values2)]),
            (
                mx.zeros((1, 8, 32, 64)),
                mx.ones((1, 8, 32, 64)),
            ),  # Standard KVCache layer
        ]

        layer_cache_types = ["CacheList", "KVCache"]

        result = ssd_cache.save_block(
            block_hash,
            cache_data,
            token_count=32,
            model_name="test",
            layer_cache_types=layer_cache_types,
        )
        assert result is True

        # Load back
        loaded = ssd_cache.load_block(block_hash)
        assert loaded is not None
        assert len(loaded) == 2

        # First layer should be List[Tuple] (CacheList)
        assert isinstance(loaded[0], list)
        assert len(loaded[0]) == 2
        assert loaded[0][0][0].shape == (1, 8, 32, 64)
        assert loaded[0][1][0].shape == (1, 4, 32, 64)

        # Second layer should be tuple (KVCache)
        assert isinstance(loaded[1], tuple)
        assert loaded[1][0].shape == (1, 8, 32, 64)

    def test_save_load_cache_list_placeholder(self, ssd_cache, mx):
        """Test saving and loading placeholder CacheList block."""
        block_hash = b"placeholder_cl_hash_"
        # Non-last block: CacheList gets standard placeholder
        cache_data = [
            (mx.zeros((1,)), mx.zeros((1,))),  # CacheList placeholder
            (mx.zeros((1, 8, 32, 64)), mx.ones((1, 8, 32, 64))),  # KVCache
        ]

        layer_cache_types = ["CacheList", "KVCache"]

        result = ssd_cache.save_block(
            block_hash,
            cache_data,
            token_count=32,
            model_name="test",
            layer_cache_types=layer_cache_types,
        )
        assert result is True

        # Load back — CacheList placeholder loads as standard (keys, values) tuple
        loaded = ssd_cache.load_block(block_hash)
        assert loaded is not None
        assert len(loaded) == 2
        # Placeholder has no sub_count, so loads as standard tuple
        assert isinstance(loaded[0], tuple)
        assert loaded[0][0].shape == (1,)

    def test_load_block_with_metadata_cache_list(self, ssd_cache, mx):
        """Test load_block_with_metadata for CacheList blocks."""
        block_hash = b"cl_metadata_test_ha_"
        sub_keys = mx.zeros((1, 8, 64, 64))
        sub_values = mx.ones((1, 8, 64, 64))

        cache_data = [
            ("__cache_list__", [(sub_keys, sub_values)]),
        ]
        layer_cache_types = ["CacheList"]
        layer_meta_states = [
            (["KVCache"], [(64,)]),  # CacheList meta_state format
        ]

        ssd_cache.save_block(
            block_hash,
            cache_data,
            token_count=64,
            model_name="test",
            layer_cache_types=layer_cache_types,
            layer_meta_states=layer_meta_states,
        )

        loaded_data, metadata = ssd_cache.load_block_with_metadata(block_hash)
        assert loaded_data is not None
        assert metadata is not None
        assert len(loaded_data) == 1
        assert isinstance(loaded_data[0], list)
        assert len(loaded_data[0]) == 1
        assert loaded_data[0][0][0].shape == (1, 8, 64, 64)
        assert metadata["layer_cache_types"] == ["CacheList"]

    def test_save_load_cache_list_with_zero_dim_values(self, ssd_cache, mx):
        """Test round-trip for CacheList where sub-cache has zero-dim values.

        This covers the deepseek_v32 / GLM-5 case where the DSA indexer
        sub-cache stores values with shape (B, 1, N, 0) — head_dim=0.
        """
        block_hash = b"zero_dim_cl_test_ha_"
        sub_keys1 = mx.zeros((1, 1, 64, 512))  # Main attention kv_latent
        sub_values1 = mx.zeros((1, 1, 64, 64))  # Main attention k_pe
        sub_keys2 = mx.zeros((1, 1, 64, 128))  # Indexer keys
        sub_values2 = mx.zeros((1, 1, 64, 0))  # Indexer values (zero head_dim)

        cache_data = [
            (
                "__cache_list__",
                [
                    (sub_keys1, sub_values1),
                    (sub_keys2, sub_values2),
                ],
            ),
        ]
        layer_cache_types = ["CacheList"]

        result = ssd_cache.save_block(
            block_hash,
            cache_data,
            token_count=64,
            model_name="test",
            layer_cache_types=layer_cache_types,
        )
        assert result is True

        # Load back and verify round-trip correctness
        loaded = ssd_cache.load_block(block_hash)
        assert loaded is not None
        assert len(loaded) == 1
        assert isinstance(loaded[0], list)
        assert len(loaded[0]) == 2

        # Sub-cache 0: normal tensors preserved
        assert loaded[0][0][0].shape == (1, 1, 64, 512)
        assert loaded[0][0][1].shape == (1, 1, 64, 64)

        # Sub-cache 1: keys normal, values zero-dim reconstructed
        assert loaded[0][1][0].shape == (1, 1, 64, 128)
        assert loaded[0][1][1].shape == (1, 1, 64, 0)

    def test_save_load_zero_dim_with_load_block_with_metadata(self, ssd_cache, mx):
        """Test load_block_with_metadata also handles zero-dim tensors."""
        block_hash = b"zero_dim_meta_test_h"
        sub_keys = mx.zeros((1, 1, 32, 128))
        sub_values = mx.zeros((1, 1, 32, 0))

        cache_data = [
            ("__cache_list__", [(sub_keys, sub_values)]),
        ]
        layer_cache_types = ["CacheList"]
        layer_meta_states = [
            (["KVCache"], [(32,)]),
        ]

        ssd_cache.save_block(
            block_hash,
            cache_data,
            token_count=32,
            model_name="test",
            layer_cache_types=layer_cache_types,
            layer_meta_states=layer_meta_states,
        )

        loaded_data, metadata = ssd_cache.load_block_with_metadata(block_hash)
        assert loaded_data is not None
        assert metadata is not None
        assert len(loaded_data) == 1
        assert isinstance(loaded_data[0], list)
        assert loaded_data[0][0][0].shape == (1, 1, 32, 128)
        assert loaded_data[0][0][1].shape == (1, 1, 32, 0)


class TestAsyncWriteAndTimeoutLoad:
    """Tests for the async write / timeout load deadlock fix.

    These tests verify:
    - save_block() returns immediately (non-blocking)
    - Pending writes are served on load (zero I/O)
    - Load timeout returns None (cache miss) instead of blocking
    - Writer thread errors clean up index entries
    - close() gracefully shuts down background threads
    """

    @pytest.fixture
    def mx(self):
        """Import MLX or skip."""
        try:
            import mlx.core as mx

            return mx
        except ImportError:
            pytest.skip("MLX not available")

    @pytest.fixture
    def ssd_cache(self, tmp_path):
        """Create a PagedSSDCacheManager for testing."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=100 * 1024**2,
        )
        yield manager
        manager.close()

    def test_save_block_non_blocking(self, ssd_cache, mx, tmp_path):
        """Verify save_block() returns immediately and file appears async."""
        block_hash = b"async_save_test_hash"
        cache_data = [
            (mx.zeros((1, 8, 64, 64)), mx.zeros((1, 8, 64, 64))) for _ in range(4)
        ]

        t0 = time.time()
        result = ssd_cache.save_block(
            block_hash=block_hash,
            cache_data=cache_data,
            token_count=64,
            model_name="test-model",
            layer_cache_types=["KVCache"] * 4,
        )
        elapsed = time.time() - t0

        assert result is True
        # save_block should return almost instantly (< 1s),
        # not wait for disk I/O
        assert elapsed < 1.0

        # Block should be in index (optimistic update)
        assert ssd_cache.has_block(block_hash)

        # Wait for background writer to finish
        import time as time_mod

        for _ in range(50):  # Wait up to 5s
            file_path = ssd_cache._get_file_path(block_hash)
            if file_path.exists():
                break
            time_mod.sleep(0.1)

        assert file_path.exists(), "File should appear after background write"

    def test_pending_writes_served_on_load(self, ssd_cache, mx):
        """Verify that a block saved then immediately loaded is served from memory."""
        block_hash = b"pending_load_test_ha"
        cache_data = [
            (mx.zeros((1, 8, 32, 64)), mx.ones((1, 8, 32, 64))) for _ in range(2)
        ]

        ssd_cache.save_block(
            block_hash=block_hash,
            cache_data=cache_data,
            token_count=32,
            model_name="test-model",
            layer_cache_types=["KVCache", "KVCache"],
        )

        # Immediately load — should come from _pending_writes, not disk
        loaded = ssd_cache.load_block(block_hash)
        assert loaded is not None
        assert len(loaded) == 2
        assert loaded[0][0].shape == (1, 8, 32, 64)
        assert loaded[0][1].shape == (1, 8, 32, 64)

    def test_pending_writes_served_on_load_with_metadata(self, ssd_cache, mx):
        """Verify load_block_with_metadata also reads from pending writes."""
        block_hash = b"pending_meta_test_ha"
        cache_data = [(mx.zeros((1, 4, 16, 32)), mx.zeros((1, 4, 16, 32)))]

        ssd_cache.save_block(
            block_hash=block_hash,
            cache_data=cache_data,
            token_count=16,
            model_name="test-model",
            layer_cache_types=["KVCache"],
            layer_meta_states=[(16,)],
        )

        loaded_data, metadata = ssd_cache.load_block_with_metadata(block_hash)
        assert loaded_data is not None
        assert metadata is not None
        assert metadata["num_layers"] == 1
        assert metadata["token_count"] == 16
        assert metadata["model_name"] == "test-model"
        assert metadata["layer_cache_types"] == ["KVCache"]

    def test_load_error_returns_none(self, ssd_cache, mx):
        """Verify that a corrupted file returns None and cleans up index."""
        block_hash = b"error_test_hash_1234"
        cache_data = [(mx.zeros((1, 8, 32, 64)), mx.zeros((1, 8, 32, 64)))]

        # Save and wait for background write to complete
        ssd_cache.save_block(
            block_hash=block_hash,
            cache_data=cache_data,
            token_count=32,
        )
        import time as time_mod

        for _ in range(50):
            with ssd_cache._pending_write_hashes_lock:
                if block_hash not in ssd_cache._pending_write_hashes:
                    break
            time_mod.sleep(0.1)

        # Remove from hot cache buffer so load goes to disk
        ssd_cache._hot_cache_remove(block_hash)

        # Mock mx.load to simulate a corrupted file
        with patch("mlx.core.load", side_effect=OSError("corrupted file")):
            loaded = ssd_cache.load_block(block_hash)
            assert loaded is None  # Should return None, not raise

        # Block should be removed from index (corrupted entry cleanup)
        assert not ssd_cache.has_block(block_hash)

    def test_load_no_executor_deadlock(self, ssd_cache, mx):
        """Regression test: _load_executor must not exist (prevents deadlock)."""
        # The old implementation used ThreadPoolExecutor(max_workers=1) which
        # caused deadlocks when mx.load() in a worker thread contested Metal
        # GPU resources with the main inference thread. Verify it's gone.
        assert not hasattr(
            ssd_cache, "_load_executor"
        ), "_load_executor should not exist — it causes Metal GPU deadlocks"

    def test_sequential_loads_no_queue_blocking(self, ssd_cache, mx):
        """Regression test: consecutive loads must not block each other."""
        import time as time_mod

        # Save 5 different blocks
        hashes = []
        for i in range(5):
            block_hash = f"seq_load_test_{i:04d}_".encode()[:20]
            hashes.append(block_hash)
            cache_data = [(mx.zeros((1, 8, 32, 64)), mx.zeros((1, 8, 32, 64)))]
            ssd_cache.save_block(block_hash, cache_data, token_count=32)

        # Wait for all pending writes to flush
        for _ in range(100):
            with ssd_cache._pending_write_hashes_lock:
                if not ssd_cache._pending_write_hashes:
                    break
            time_mod.sleep(0.1)

        # Load all 5 blocks sequentially — should complete quickly
        t0 = time_mod.time()
        for block_hash in hashes:
            loaded = ssd_cache.load_block(block_hash)
            assert loaded is not None, f"Failed to load {block_hash!r}"
            assert len(loaded) == 1
        elapsed = time_mod.time() - t0

        # 5 loads from SSD should complete in well under 5s
        # (each ~2ms read + reconstruction)
        assert (
            elapsed < 5.0
        ), f"Sequential loads took {elapsed:.1f}s — possible queue blocking"

    def test_writer_error_handling(self, ssd_cache, mx):
        """Verify that background writer errors clean up the index."""
        block_hash = b"writer_error_test_ha"
        cache_data = [(mx.zeros((1, 4, 16, 32)), mx.zeros((1, 4, 16, 32)))]

        # Patch _write_safetensors_no_mx to simulate disk error in background writer
        import time as time_mod

        with patch(
            "omlx.cache.paged_ssd_cache._write_safetensors_no_mx",
            side_effect=OSError("Disk full"),
        ):
            result = ssd_cache.save_block(
                block_hash=block_hash,
                cache_data=cache_data,
                token_count=16,
            )
            # save_block() succeeds (bytes extracted, queued for background write)
            assert result is True

            # Wait for background writer to process and fail
            for _ in range(50):
                if ssd_cache._write_queue.empty():
                    break
                time_mod.sleep(0.05)
            time_mod.sleep(0.1)

        # Background writer should have removed the block from index on error
        assert not ssd_cache.has_block(block_hash)
        # And from pending write hashes
        with ssd_cache._pending_write_hashes_lock:
            assert block_hash not in ssd_cache._pending_write_hashes

    def test_writer_enospc_logs_disk_full(self, ssd_cache, mx, caplog):
        """ENOSPC errors should log 'disk full' warning, not generic error."""
        block_hash = b"enospc_test_hash_123"
        cache_data = [(mx.zeros((1, 4, 16, 32)), mx.zeros((1, 4, 16, 32)))]

        enospc = OSError("No space left on device")
        enospc.errno = errno.ENOSPC

        import time as time_mod

        with (
            patch(
                "omlx.cache.paged_ssd_cache._write_safetensors_no_mx",
                side_effect=enospc,
            ),
            caplog.at_level(logging.WARNING),
        ):
            ssd_cache.save_block(
                block_hash=block_hash,
                cache_data=cache_data,
                token_count=16,
            )
            for _ in range(50):
                if ssd_cache._write_queue.empty():
                    break
                time_mod.sleep(0.05)
            time_mod.sleep(0.1)

        assert "SSD cache disk full" in caplog.text

    def test_graceful_shutdown(self, tmp_path, mx):
        """Verify close() stops the writer thread."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "shutdown_cache",
            max_size_bytes=100 * 1024**2,
        )

        # Save a block to ensure writer is active
        block_hash = b"shutdown_test_hash_1"
        cache_data = [(mx.zeros((1, 4, 16, 32)), mx.zeros((1, 4, 16, 32)))]
        manager.save_block(block_hash, cache_data, 16)

        # Close should stop the writer thread
        manager.close()

        assert not manager._writer_thread.is_alive()

    def test_save_existing_block_still_touches(self, ssd_cache, mx):
        """Verify saving an existing block just touches LRU (unchanged behavior)."""
        block_hash = b"touch_existing_test_"
        cache_data = [(mx.zeros((1, 8, 32, 64)), mx.zeros((1, 8, 32, 64)))]

        ssd_cache.save_block(block_hash, cache_data, 32)
        initial_saves = ssd_cache._stats["saves"]

        # Second save should just touch, not re-enqueue
        ssd_cache.save_block(block_hash, cache_data, 32)
        assert ssd_cache._stats["saves"] == initial_saves
        assert ssd_cache._stats["hits"] >= 1

    def test_save_and_load_round_trip_after_flush(self, ssd_cache, mx):
        """Verify full round-trip: save -> flush -> load from disk."""
        import time as time_mod

        block_hash = b"round_trip_flush_tes"
        cache_data = [
            (mx.zeros((1, 8, 64, 64)), mx.ones((1, 8, 64, 64))) for _ in range(4)
        ]

        ssd_cache.save_block(
            block_hash=block_hash,
            cache_data=cache_data,
            token_count=64,
            model_name="test-model",
            layer_cache_types=["KVCache"] * 4,
        )

        # Wait for background write to complete
        for _ in range(50):
            with ssd_cache._pending_write_hashes_lock:
                if block_hash not in ssd_cache._pending_write_hashes:
                    break
            time_mod.sleep(0.1)

        # Remove from hot cache buffer so load goes to disk
        ssd_cache._hot_cache_remove(block_hash)

        # Now load should come from disk, not pending writes
        loaded = ssd_cache.load_block(block_hash)
        assert loaded is not None
        assert len(loaded) == 4
        for keys, values in loaded:
            assert keys.shape == (1, 8, 64, 64)
            assert values.shape == (1, 8, 64, 64)


# =============================================================================
# Async Background Write Tests
# =============================================================================


@pytest.mark.skipif(not _has_mlx(), reason="MLX not available")
class TestAsyncBackgroundWrite:
    """Tests for the async background write pipeline (no-mx safetensors)."""

    @pytest.fixture
    def mx(self):
        import mlx.core as mx

        return mx

    def test_extract_and_restore_float32(self, mx):
        """Round-trip test for float32 tensors."""
        original = mx.random.normal((2, 4, 8))
        mx.eval(original)
        raw, dtype_str, shape = _extract_tensor_bytes(original)
        assert dtype_str == "F32"
        assert shape == [2, 4, 8]
        restored = _restore_tensor_from_bytes(raw, dtype_str, shape)
        assert restored.dtype == mx.float32
        assert restored.shape == (2, 4, 8)
        assert mx.allclose(original, restored).item()

    def test_extract_and_restore_float16(self, mx):
        """Round-trip test for float16 tensors."""
        original = mx.random.normal((3, 5)).astype(mx.float16)
        mx.eval(original)
        raw, dtype_str, shape = _extract_tensor_bytes(original)
        assert dtype_str == "F16"
        restored = _restore_tensor_from_bytes(raw, dtype_str, shape)
        assert restored.dtype == mx.float16
        assert mx.allclose(original, restored).item()

    def test_extract_and_restore_bfloat16(self, mx):
        """Round-trip test for bfloat16 tensors (the key dtype for this feature)."""
        original = mx.random.normal((4, 8, 16)).astype(mx.bfloat16)
        mx.eval(original)
        raw, dtype_str, shape = _extract_tensor_bytes(original)
        assert dtype_str == "BF16"
        assert shape == [4, 8, 16]
        restored = _restore_tensor_from_bytes(raw, dtype_str, shape)
        assert restored.dtype == mx.bfloat16
        assert restored.shape == (4, 8, 16)
        # Compare as float32 to avoid bfloat16 precision issues
        assert mx.allclose(
            original.astype(mx.float32), restored.astype(mx.float32)
        ).item()

    def test_extract_and_restore_int_types(self, mx):
        """Round-trip test for integer dtypes."""
        for mx_dtype, st_str in [
            (mx.int8, "I8"),
            (mx.int32, "I32"),
            (mx.uint8, "U8"),
        ]:
            original = mx.array([1, 2, 3, 4], dtype=mx_dtype)
            mx.eval(original)
            raw, dtype_str, shape = _extract_tensor_bytes(original)
            assert dtype_str == st_str
            restored = _restore_tensor_from_bytes(raw, dtype_str, shape)
            assert restored.dtype == mx_dtype
            assert mx.array_equal(original, restored).item()

    def test_extract_materializes_lazy_slice(self, mx):
        """_extract_tensor_bytes handles lazy block slices."""
        base = mx.arange(1 * 2 * 16 * 4, dtype=mx.float32).reshape(1, 2, 16, 4)
        mx.eval(base)
        lazy_slice = base[:, :, 3:11, :]

        raw, dtype_str, shape = _extract_tensor_bytes(lazy_slice)

        restored = _restore_tensor_from_bytes(raw, dtype_str, shape)
        expected = base[:, :, 3:11, :]
        mx.eval(expected)
        assert dtype_str == "F32"
        assert shape == [1, 2, 8, 4]
        assert mx.allclose(expected, restored).item()

    def test_extract_materializes_lazy_bfloat16_slice(self, mx):
        """_extract_tensor_bytes handles lazy bf16 slices and uint16 views."""
        base = mx.arange(1 * 2 * 12 * 4, dtype=mx.float32).reshape(1, 2, 12, 4)
        base = base.astype(mx.bfloat16)
        mx.eval(base)
        lazy_slice = base[:, :, 2:10, :]

        raw, dtype_str, shape = _extract_tensor_bytes(lazy_slice)

        restored = _restore_tensor_from_bytes(raw, dtype_str, shape)
        expected = base[:, :, 2:10, :]
        mx.eval(expected)
        assert dtype_str == "BF16"
        assert shape == [1, 2, 8, 4]
        assert restored.dtype == mx.bfloat16
        assert mx.allclose(
            expected.astype(mx.float32), restored.astype(mx.float32)
        ).item()

    def test_extract_materializes_lazy_clone(self, mx):
        """_extract_tensor_bytes handles block-like lazy clone/copy tensors."""
        base = mx.arange(1 * 2 * 16 * 4, dtype=mx.float32).reshape(1, 2, 16, 4)
        mx.eval(base)
        tensor = base[:, :, 4:12, :]
        if hasattr(mx, "copy"):
            cloned = mx.copy(tensor)
        elif hasattr(tensor, "copy"):
            cloned = tensor.copy()
        else:
            cloned = mx.array(tensor)

        raw, dtype_str, shape = _extract_tensor_bytes(cloned)

        restored = _restore_tensor_from_bytes(raw, dtype_str, shape)
        expected = base[:, :, 4:12, :]
        mx.eval(expected)
        assert dtype_str == "F32"
        assert shape == [1, 2, 8, 4]
        assert mx.allclose(expected, restored).item()

    def test_write_safetensors_no_mx_roundtrip(self, mx, tmp_path):
        """Write safetensors without mx API, then load with mx.load()."""
        t1 = mx.random.normal((2, 3, 4))
        t2 = mx.ones((5,), dtype=mx.float16)
        mx.eval(t1, t2)

        tensors_raw = {
            "tensor_a": _extract_tensor_bytes(t1),
            "tensor_b": _extract_tensor_bytes(t2),
        }
        metadata = {"test_key": "test_value", "block_hash": "abc123"}

        out_path = str(tmp_path / "test.safetensors")
        file_size = _write_safetensors_no_mx(out_path, tensors_raw, metadata)
        assert file_size > 0

        # Load with mx.load and verify
        loaded_arrays, loaded_meta = mx.load(out_path, return_metadata=True)
        assert "tensor_a" in loaded_arrays
        assert "tensor_b" in loaded_arrays
        assert loaded_meta["test_key"] == "test_value"
        assert loaded_meta["block_hash"] == "abc123"
        assert mx.allclose(t1, loaded_arrays["tensor_a"]).item()
        assert mx.allclose(t2, loaded_arrays["tensor_b"]).item()

    def test_write_safetensors_bfloat16_roundtrip(self, mx, tmp_path):
        """Verify bfloat16 safetensors file is loadable by mx.load."""
        original = mx.random.normal((8, 16, 32)).astype(mx.bfloat16)
        mx.eval(original)

        tensors_raw = {"kv_cache": _extract_tensor_bytes(original)}
        out_path = str(tmp_path / "bf16_test.safetensors")
        _write_safetensors_no_mx(out_path, tensors_raw)

        loaded, _ = mx.load(out_path, return_metadata=True)
        assert loaded["kv_cache"].dtype == mx.bfloat16
        assert loaded["kv_cache"].shape == (8, 16, 32)
        assert mx.allclose(
            original.astype(mx.float32),
            loaded["kv_cache"].astype(mx.float32),
        ).item()

    def test_save_block_uses_background_write(self, tmp_path, mx):
        """Verify save_block enqueues bytes for background writer (no mx.save_safetensors)."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "async_test",
            max_size_bytes=100 * 1024**2,
        )

        block_hash = b"async_write_test_hsh"
        cache_data = [(mx.ones((1, 4, 16, 32)), mx.zeros((1, 4, 16, 32)))]

        # Patch mx.save_safetensors to ensure it's NOT called
        with patch("mlx.core.save_safetensors") as mock_save:
            result = manager.save_block(
                block_hash=block_hash,
                cache_data=cache_data,
                token_count=16,
            )
            assert result is True
            # mx.save_safetensors should NOT be called (we use _write_safetensors_no_mx)
            mock_save.assert_not_called()

        # Hot cache buffer should store tensors_raw (bytes), not arrays (mx.array)
        with manager._hot_cache_lock:
            pending = manager._hot_cache.get(block_hash)
        assert pending is not None
        assert "tensors_raw" in pending
        assert "arrays" not in pending  # Old key should not exist

        # Wait for background write and verify file exists
        for _ in range(50):
            file_path = manager._get_file_path(block_hash)
            if file_path.exists():
                break
            time.sleep(0.05)
        assert file_path.exists()

        # Verify file is loadable by mx.load. V3 stores state elements as
        # ``layer_{i}_state_{k}`` keys with a ``layer_{i}_state_count`` meta
        # entry, polyfilled from V2 ``(keys, values)`` 2-tuples on save.
        loaded, meta = mx.load(str(file_path), return_metadata=True)
        assert "layer_0_state_0" in loaded
        assert "layer_0_state_1" in loaded
        assert meta.get("layer_0_state_count") == "2"
        assert meta["block_hash"] == block_hash.hex()

        manager.close()

    def test_pending_writes_bytes_readback(self, tmp_path, mx):
        """Verify load_block can restore mx.arrays from bytes-based pending_writes."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "readback_test",
            max_size_bytes=100 * 1024**2,
        )

        block_hash = b"readback_test_hash__"
        original_keys = mx.random.normal((1, 8, 32, 64))
        original_values = mx.random.normal((1, 8, 32, 64))
        mx.eval(original_keys, original_values)
        cache_data = [(original_keys, original_values)]

        manager.save_block(
            block_hash=block_hash,
            cache_data=cache_data,
            token_count=32,
        )

        # Load immediately from pending_writes (before background write completes)
        loaded = manager.load_block(block_hash)
        assert loaded is not None
        assert len(loaded) == 1
        keys, values = loaded[0]
        assert mx.allclose(original_keys, keys).item()
        assert mx.allclose(original_values, values).item()

        manager.close()

    def test_index_update_file_size(self):
        """Verify PagedSSDCacheIndex.update_file_size works correctly."""
        index = PagedSSDCacheIndex(max_size_bytes=1000)
        block_hash = b"size_update_test____"
        metadata = PagedSSDBlockMetadata(
            block_hash=block_hash,
            file_path=Path("/tmp/test.safetensors"),
            file_size=100,
            token_count=16,
            created_at=time.time(),
            last_access=time.time(),
            num_layers=1,
        )
        index.add(metadata)
        assert index.total_size == 100

        # Update to actual size
        index.update_file_size(block_hash, 150)
        assert index.total_size == 150

        # Non-existent hash should be no-op
        index.update_file_size(b"nonexistent_hash____", 999)
        assert index.total_size == 150


class TestEffectiveMaxSize:
    """Tests for dynamic effective max size based on disk free space."""

    def _make_disk_usage(self, total: int, used: int, free: int):
        """Create a mock disk_usage result."""
        return shutil._ntuple_diskusage(total, used, free)

    def test_effective_max_size_disk_sufficient(self, tmp_path: Path):
        """When disk has plenty of free space, effective = configured max."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=100 * 1024**3,  # 100GB configured
        )

        # Mock: 500GB free, cache is empty (0 bytes)
        mock_usage = self._make_disk_usage(
            total=1000 * 1024**3, used=500 * 1024**3, free=500 * 1024**3
        )
        with patch("shutil.disk_usage", return_value=mock_usage):
            effective = manager._get_effective_max_size()

        # disk_available = 0 + 500GB = 500GB, disk_limit = 495GB
        # effective = min(100GB, 495GB) = 100GB
        assert effective == 100 * 1024**3

    def test_effective_max_size_disk_low(self, tmp_path: Path):
        """When disk is low, effective shrinks below configured max."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=110 * 1024**3,  # 110GB configured
        )

        # Simulate: cache currently has 10GB, disk free is 90GB
        # So disk_available = 10GB + 90GB = 100GB
        manager._index._total_size = 10 * 1024**3

        mock_usage = self._make_disk_usage(
            total=500 * 1024**3, used=410 * 1024**3, free=90 * 1024**3
        )
        with patch("shutil.disk_usage", return_value=mock_usage):
            effective = manager._get_effective_max_size()

        # disk_limit = int(100GB * 0.99) = 99GB
        # effective = min(110GB, 99GB) = 99GB
        expected = int(100 * 1024**3 * 0.99)
        assert effective == expected

    def test_effective_max_size_oserror_fallback(self, tmp_path: Path):
        """When disk_usage fails, fall back to configured max."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=50 * 1024**3,
        )

        with patch("shutil.disk_usage", side_effect=OSError("disk error")):
            effective = manager._get_effective_max_size()

        assert effective == 50 * 1024**3

    def test_effective_max_size_cache_30s(self, tmp_path: Path):
        """disk_usage result is cached for 30 seconds."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=100 * 1024**3,
        )

        mock_usage = self._make_disk_usage(
            total=1000 * 1024**3, used=500 * 1024**3, free=500 * 1024**3
        )
        with patch("shutil.disk_usage", return_value=mock_usage) as mock_du:
            # First call — should invoke disk_usage
            manager._get_effective_max_size()
            assert mock_du.call_count == 1

            # Second call within 30s — should use cache
            manager._get_effective_max_size()
            assert mock_du.call_count == 1

            # Expire cache by rewinding timestamp
            manager._disk_usage_cache_time -= 31.0

            # Third call — should invoke disk_usage again
            manager._get_effective_max_size()
            assert mock_du.call_count == 2

    def test_utilization_never_exceeds_1(self, tmp_path: Path):
        """Utilization should never exceed 1.0 with effective max size."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=100 * 1024**3,
        )

        # Simulate: cache has 50GB, but disk only has 10GB free
        # So disk_available = 50GB + 10GB = 60GB, disk_limit = ~59.4GB
        manager._index._total_size = 50 * 1024**3

        mock_usage = self._make_disk_usage(
            total=200 * 1024**3, used=190 * 1024**3, free=10 * 1024**3
        )
        with patch("shutil.disk_usage", return_value=mock_usage):
            stats = manager.get_stats_dict()

        assert stats["utilization"] <= 1.0
        assert stats["max_size"] < stats["configured_max_size"]

    def test_effective_max_size_uses_tracked_cache_not_filesystem_total(
        self, tmp_path: Path
    ):
        """Disk pressure depends on tracked cache bytes plus free space."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=100 * 1024**3,
        )
        manager._index._total_size = 50 * 1024**3
        manager._incompatible_index._total_size = 10 * 1024**3

        mock_usage = self._make_disk_usage(
            total=200 * 1024**3,
            used=190 * 1024**3,
            free=10 * 1024**3,
        )
        with patch("shutil.disk_usage", return_value=mock_usage):
            effective = manager._get_effective_max_size()

        expected = int(70 * 1024**3 * 0.99)
        assert effective == expected
        assert effective < manager.configured_max_size

    def test_stats_includes_effective_and_configured(self, tmp_path: Path):
        """Stats should include both effective and configured max sizes."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=100 * 1024**3,
        )

        mock_usage = self._make_disk_usage(
            total=500 * 1024**3, used=450 * 1024**3, free=50 * 1024**3
        )
        with patch("shutil.disk_usage", return_value=mock_usage):
            stats_dict = manager.get_stats_dict()
            stats_obj = manager.get_stats()

        # Dict format
        assert "configured_max_size" in stats_dict
        assert stats_dict["configured_max_size"] == 100 * 1024**3

        # Dataclass format
        assert stats_obj.configured_max_size_bytes == 100 * 1024**3
        assert stats_obj.max_size_bytes > 0

    def test_max_size_property_returns_effective(self, tmp_path: Path):
        """max_size property should return effective (not configured) value."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=200 * 1024**3,
        )

        # disk_available = 0 + 50GB = 50GB, disk_limit = ~49.5GB
        mock_usage = self._make_disk_usage(
            total=500 * 1024**3, used=450 * 1024**3, free=50 * 1024**3
        )
        with patch("shutil.disk_usage", return_value=mock_usage):
            assert manager.max_size < 200 * 1024**3
            assert manager.configured_max_size == 200 * 1024**3

    def test_oserror_fallback_logs_warning(self, tmp_path: Path, caplog):
        """disk_usage failure should log a warning, not fail silently."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=50 * 1024**3,
        )
        # Expire cache so next call hits disk_usage
        manager._disk_usage_cache_time -= 31.0

        with (
            patch("shutil.disk_usage", side_effect=OSError("mount gone")),
            caplog.at_level(logging.WARNING),
        ):
            effective = manager._get_effective_max_size()

        assert effective == 50 * 1024**3
        assert "Failed to check disk usage" in caplog.text

    def test_disk_pressure_warning(self, tmp_path: Path, caplog):
        """Warn when effective max drops below 10% of configured max."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=100 * 1024**3,
        )

        # Simulate nearly full disk: only 5GB free, cache has 0 bytes
        mock_usage = self._make_disk_usage(
            total=500 * 1024**3, used=495 * 1024**3, free=5 * 1024**3
        )
        with (
            patch("shutil.disk_usage", return_value=mock_usage),
            caplog.at_level(logging.WARNING),
        ):
            manager._enforce_size_limit_for_new_block()

        assert "disk pressure" in caplog.text
        assert "disk nearly full" in caplog.text


class TestPreloadMatchedBlocks:
    """Tests for parallel block preloading into hot cache."""

    @pytest.fixture
    def mx(self):
        try:
            import mlx.core as mx

            return mx
        except ImportError:
            pytest.skip("MLX not available")

    @pytest.fixture
    def manager_with_hot_cache(self, tmp_path, mx):
        """Create a manager with hot cache enabled."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
            hot_cache_max_bytes=512 * 1024**2,
        )
        yield manager
        manager.close()

    def _save_test_blocks(
        self,
        manager,
        mx,
        count=4,
        layers=2,
        hot_cache_max_bytes=512 * 1024**2,
        hot_cache_budget=None,
    ):
        """Save test blocks and flush them to SSD (not hot cache)."""
        hashes = []
        for i in range(count):
            block_hash = f"preload_test_block_{i:04d}".encode()
            cache_data = [
                (
                    mx.zeros((1, 4, 64, 64)),
                    mx.zeros((1, 4, 64, 64)),
                )
                for _ in range(layers)
            ]
            manager.save_block(
                block_hash=block_hash,
                cache_data=cache_data,
                token_count=64,
                model_name="test-model",
                layer_cache_types=["KVCache"] * layers,
            )
            hashes.append(block_hash)

        # Flush writer to ensure blocks are on SSD
        manager.close()

        # Re-open manager (cold start — hot cache is empty)
        new_manager = PagedSSDCacheManager(
            cache_dir=manager._cache_dir,
            max_size_bytes=1024**3,
            hot_cache_max_bytes=hot_cache_max_bytes,
            hot_cache_budget=hot_cache_budget,
        )
        return new_manager, hashes

    def test_preload_promotes_to_hot_cache(self, tmp_path, mx):
        """After preload, blocks are found in hot cache."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
            hot_cache_max_bytes=512 * 1024**2,
        )
        manager2, hashes = self._save_test_blocks(manager, mx, count=4)

        # Verify blocks are NOT in hot cache before preload
        for h in hashes:
            assert manager2._hot_cache_get(h) is None

        # Preload
        loaded = manager2.preload_matched_blocks(hashes)
        assert loaded == 4

        # Verify blocks ARE in hot cache after preload
        for h in hashes:
            assert manager2._hot_cache_get(h) is not None

        manager2.close()

    def test_clean_promoted_block_eviction_skips_ssd_write(self, tmp_path, mx):
        """Blocks loaded from SSD are clean and should not be re-written."""
        entry_size = 2 * 2 * 1 * 4 * 64 * 64 * 4
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
            hot_cache_max_bytes=512 * 1024**2,
        )
        manager2, hashes = self._save_test_blocks(
            manager,
            mx,
            count=2,
            hot_cache_max_bytes=entry_size + 1024,
        )

        try:
            assert manager2.load_block(hashes[0]) is not None
            first_entry = manager2._hot_cache_get(hashes[0])
            assert first_entry is not None
            assert first_entry["dirty"] is False

            with patch.object(manager2, "_enqueue_ssd_write") as enqueue_write:
                assert manager2.load_block(hashes[1]) is not None
                enqueue_write.assert_not_called()

            assert manager2._hot_cache_get(hashes[0]) is None
            assert manager2._hot_cache_get(hashes[1]) is not None
        finally:
            manager2.close()

    def test_preload_partial_failure(self, tmp_path, mx):
        """If one block file is missing, others still load."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
            hot_cache_max_bytes=512 * 1024**2,
        )
        manager2, hashes = self._save_test_blocks(manager, mx, count=5)

        # Delete one block file from SSD to simulate failure
        metadata = manager2._index.get(hashes[1])
        metadata.file_path.unlink()

        loaded = manager2.preload_matched_blocks(hashes)

        # 4 of 5 should succeed (1 deleted)
        assert loaded == 4
        assert manager2._hot_cache_get(hashes[0]) is not None
        assert manager2._hot_cache_get(hashes[1]) is None  # deleted file
        assert manager2._hot_cache_get(hashes[2]) is not None

        manager2.close()

    def test_preload_skips_hot_cache_blocks(self, tmp_path, mx):
        """Blocks already in hot cache are not re-loaded."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
            hot_cache_max_bytes=512 * 1024**2,
        )
        manager2, hashes = self._save_test_blocks(manager, mx, count=5)

        # Load one block into hot cache manually
        manager2.load_block(hashes[0])
        assert manager2._hot_cache_get(hashes[0]) is not None
        promotions_before = manager2._stats["hot_cache_promotions"]

        # Preload all — should only load the 4 cold blocks
        loaded = manager2.preload_matched_blocks(hashes)
        assert loaded == 4

        # Promotion count should increase by exactly 4 (not 5)
        assert manager2._stats["hot_cache_promotions"] == promotions_before + 4

        manager2.close()

    def test_preload_unknown_hashes_ignored(self, tmp_path, mx):
        """Hashes not in the SSD index are silently skipped."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
            hot_cache_max_bytes=512 * 1024**2,
        )
        manager2, hashes = self._save_test_blocks(manager, mx, count=5)

        all_hashes = hashes + [b"nonexistent_hash_01", b"nonexistent_hash_02"]
        loaded = manager2.preload_matched_blocks(all_hashes)
        assert loaded == 5  # only the real blocks

        manager2.close()

    def test_preload_noop_without_hot_cache(self, tmp_path, mx):
        """Preload returns 0 when hot cache is disabled."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
            hot_cache_max_bytes=0,  # hot cache disabled
        )
        block_hash = b"preload_no_hot_test"
        cache_data = [(mx.zeros((1, 4, 32, 64)), mx.zeros((1, 4, 32, 64)))]
        manager.save_block(block_hash, cache_data, 32, layer_cache_types=["KVCache"])
        manager.close()

        manager2 = PagedSSDCacheManager(
            cache_dir=manager._cache_dir,
            max_size_bytes=1024**3,
            hot_cache_max_bytes=0,
        )
        loaded = manager2.preload_matched_blocks([block_hash])
        assert loaded == 0

        manager2.close()

    def test_preload_skips_when_hot_cache_full(self, tmp_path, mx):
        """Preload returns 0 when hot cache has no remaining capacity."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
            hot_cache_max_bytes=1024,  # tiny hot cache
        )
        manager2, hashes = self._save_test_blocks(manager, mx, count=2)

        # Fill hot cache to capacity
        manager2._hot_cache_total_bytes = manager2._hot_cache_max_bytes

        loaded = manager2.preload_matched_blocks(hashes)
        assert loaded == 0

        manager2.close()

    def test_preload_respects_available_hot_cache_capacity(self, tmp_path, mx):
        """Preload must not load more cold blocks than remaining hot-cache bytes."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
            hot_cache_max_bytes=512 * 1024**2,
        )
        manager2, hashes = self._save_test_blocks(manager, mx, count=5)

        try:
            file_sizes = [manager2._index.get(h).file_size for h in hashes]
            manager2._hot_cache_max_bytes = sum(file_sizes[:4])

            loaded = manager2.preload_matched_blocks(hashes)

            assert loaded == 4
            for h in hashes[:4]:
                assert manager2._hot_cache_get(h) is not None
            assert manager2._hot_cache_get(hashes[4]) is None
        finally:
            manager2.close()

    def test_preload_skips_when_shared_hot_cache_budget_full(self, tmp_path, mx):
        """Preload uses remaining shared budget, not only local hot cache bytes."""
        budget = SharedHotCacheBudget(1024)
        filler = PagedSSDCacheManager(
            cache_dir=tmp_path / "budget_filler",
            max_size_bytes=1024**3,
            hot_cache_max_bytes=budget.max_bytes,
            hot_cache_only=True,
            hot_cache_budget=budget,
        )
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
            hot_cache_max_bytes=0,
        )
        manager2 = None

        try:
            filler._hot_cache_put(
                b"preload_budget_filler",
                {
                    "tensors_raw": {
                        "layer_0_keys": (bytes(1024), "uint8", [1024]),
                    },
                    "file_metadata": {},
                    "num_layers": 1,
                    "layer_cache_types": ["KVCache"],
                    "block_metadata": None,
                },
            )
            assert budget.remaining_bytes == 0

            manager2, hashes = self._save_test_blocks(
                manager,
                mx,
                count=4,
                hot_cache_max_bytes=budget.max_bytes,
                hot_cache_budget=budget,
            )

            loaded = manager2.preload_matched_blocks(hashes)
            assert loaded == 0
            assert budget.total_bytes == budget.max_bytes
            for h in hashes:
                assert manager2._hot_cache_get(h) is None
        finally:
            if manager2 is not None:
                manager2.close()
            else:
                manager.close()
            filler.close()

    def test_preload_empty_list(self, manager_with_hot_cache):
        """Empty hash list returns 0 immediately."""
        loaded = manager_with_hot_cache.preload_matched_blocks([])
        assert loaded == 0

    def test_preload_skips_below_threshold(self, tmp_path, mx):
        """Preload skips when fewer than 4 cold blocks need loading."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
            hot_cache_max_bytes=512 * 1024**2,
        )
        manager2, hashes = self._save_test_blocks(manager, mx, count=3)

        loaded = manager2.preload_matched_blocks(hashes)
        assert loaded == 0
        for bh in hashes:
            assert manager2._hot_cache_get(bh) is None

        manager2.close()

    def test_preload_updates_stats(self, tmp_path, mx):
        """Preload increments preload-specific stats counters."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
            hot_cache_max_bytes=512 * 1024**2,
        )
        manager2, hashes = self._save_test_blocks(manager, mx, count=4)

        manager2.preload_matched_blocks(hashes)

        assert manager2._stats["preload_blocks_loaded"] == 4
        assert manager2._stats["preload_calls"] == 1
        assert manager2._stats["preload_time_ms"] > 0

        manager2.close()

    def test_preloaded_blocks_load_correctly(self, tmp_path, mx):
        """After preload, load_block returns correct data from hot cache."""
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
            hot_cache_max_bytes=512 * 1024**2,
        )
        manager2, hashes = self._save_test_blocks(manager, mx, count=5, layers=3)

        # Preload blocks into hot cache
        manager2.preload_matched_blocks(hashes)

        # Now load_block should hit hot cache
        hot_hits_before = manager2._stats["hot_cache_hits"]
        for h in hashes:
            data = manager2.load_block(h)
            assert data is not None
            assert len(data) == 3  # 3 layers
            for keys, values in data:
                assert keys.shape == (1, 4, 64, 64)
                assert values.shape == (1, 4, 64, 64)

        # All loads should be hot cache hits (not SSD reads)
        assert manager2._stats["hot_cache_hits"] == hot_hits_before + 5

        manager2.close()

    def test_concurrent_preload_and_load(self, tmp_path, mx):
        """Preload and load_block don't race on hot cache."""
        import threading

        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
            hot_cache_max_bytes=512 * 1024**2,
        )
        manager2, hashes = self._save_test_blocks(manager, mx, count=8)

        results = {"preload": None, "loads": []}
        errors = []

        def do_preload():
            try:
                results["preload"] = manager2.preload_matched_blocks(hashes)
            except Exception as e:
                errors.append(f"preload: {e}")

        def do_loads():
            try:
                for h in hashes:
                    data = manager2.load_block(h)
                    results["loads"].append(data is not None)
            except Exception as e:
                errors.append(f"load: {e}")

        t1 = threading.Thread(target=do_preload)
        t2 = threading.Thread(target=do_loads)
        t1.start()
        t2.start()
        t1.join(timeout=30)
        t2.join(timeout=30)

        assert not errors, f"Concurrent errors: {errors}"
        assert not t1.is_alive(), "Preload thread hung"
        assert not t2.is_alive(), "Load thread hung"

        manager2.close()


class TestPreloadBlocks:
    """Tests for BlockAwarePrefixCache.preload_blocks()."""

    @pytest.fixture
    def mx(self):
        try:
            import mlx.core as mx

            return mx
        except ImportError:
            pytest.skip("MLX not available")

    def test_preload_blocks_calls_ssd_preload(self, tmp_path, mx):
        """preload_blocks extracts hashes from BlockTable and calls SSD preload."""
        from unittest.mock import MagicMock

        from omlx.cache.paged_cache import PagedCacheManager

        # Set up real SSD manager with blocks
        ssd_manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
            hot_cache_max_bytes=256 * 1024**2,
        )

        hashes = []
        for i in range(5):
            bh = f"preload_blocks_test_{i:04d}".encode()
            cache_data = [(mx.zeros((1, 4, 32, 64)), mx.zeros((1, 4, 32, 64)))]
            ssd_manager.save_block(bh, cache_data, 32, layer_cache_types=["KVCache"])
            hashes.append(bh)
        ssd_manager.close()

        # Re-open cold
        ssd_manager2 = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1024**3,
            hot_cache_max_bytes=256 * 1024**2,
        )

        # Set up paged cache with allocated blocks
        paged_cache = PagedCacheManager(block_size=256, max_blocks=100)
        block_ids = []
        for bh in hashes:
            block = paged_cache.allocate_block()
            block.block_hash = bh
            block.token_count = 32
            block_ids.append(block.block_id)

        # Create BlockAwarePrefixCache
        from omlx.cache.prefix_cache import BlockAwarePrefixCache, BlockTable

        model = MagicMock()
        prefix_cache = BlockAwarePrefixCache(model, paged_cache, ssd_manager2)

        # Create a BlockTable
        bt = BlockTable(
            request_id="test-req",
            block_ids=block_ids,
            num_tokens=32 * len(block_ids),
        )

        # Call preload_blocks
        loaded = prefix_cache.preload_blocks(bt)
        assert loaded == 5

        # Verify blocks are in hot cache
        for bh in hashes:
            assert ssd_manager2._hot_cache_get(bh) is not None

        ssd_manager2.close()


class TestComputeMaxPendingWrites:
    """Pin the block-size-aware cap formula.

    The pending-writes queue holds raw KV bytes that the background
    SSD writer hasn't drained yet — so the cap must scale by bytes per
    slot, not just host RAM. Cap should shrink when each block is
    expensive (large block_size or fat KV) and grow when each block is
    cheap, while the hard byte budget bounds the soft floor.
    """

    def test_soft_floor_applies_when_hard_budget_allows(self):
        """The soft floor gives large-block workloads burst headroom
        when the byte hard cap still has room."""
        from omlx.cache.paged_ssd_cache import _compute_max_pending_writes

        def fake_sysconf(name):
            if name == "SC_PAGE_SIZE":
                return 4096
            if name == "SC_PHYS_PAGES":
                return (64 * 1024**3) // 4096
            raise ValueError(name)

        with patch("omlx.cache.paged_ssd_cache.os.sysconf", side_effect=fake_sysconf):
            cap = _compute_max_pending_writes(
                block_size_tokens=2048,
                kv_bytes_per_token=200_000,
            )
        assert cap == 32

    def test_hard_budget_bounds_soft_floor(self):
        """The soft floor must not force the pending pool above 30%
        of host RAM for very expensive blocks."""
        from omlx.cache.paged_ssd_cache import _compute_max_pending_writes

        def fake_sysconf(name):
            if name == "SC_PAGE_SIZE":
                return 4096
            if name == "SC_PHYS_PAGES":
                return (64 * 1024**3) // 4096
            raise ValueError(name)

        with patch("omlx.cache.paged_ssd_cache.os.sysconf", side_effect=fake_sysconf):
            cap = _compute_max_pending_writes(
                block_size_tokens=2048,
                kv_bytes_per_token=500_000,
            )
        assert cap == 20

    def test_ceiling_clamps_high_end(self):
        """A 512 GB host with tiny blocks still tops out at 256
        rather than pinning thousands of in-flight writes."""
        from omlx.cache.paged_ssd_cache import _compute_max_pending_writes

        # Tiny per-slot cost → math would produce a huge number;
        # the ceiling must hold at 256.
        cap = _compute_max_pending_writes(
            block_size_tokens=1,
            kv_bytes_per_token=1,
        )
        assert cap == 256

    def test_larger_block_shrinks_cap(self):
        """Doubling block_size_tokens halves the cap target (until
        floor/ceiling clamp)."""
        from omlx.cache.paged_ssd_cache import _compute_max_pending_writes

        cap_256 = _compute_max_pending_writes(block_size_tokens=256)
        cap_1024 = _compute_max_pending_writes(block_size_tokens=1024)

        # 4× block size → ¼ the target. Both should be in the
        # in-band range (above floor, below ceiling) for typical
        # Macs; if the smaller one is at the ceiling the test
        # degenerates but the ordering still holds.
        assert cap_1024 <= cap_256, (
            "Larger blocks must yield smaller cap so worst-case "
            "pinned bytes stay bounded"
        )

    def test_larger_kv_bytes_shrinks_cap(self):
        """Heavier per-token KV (bigger model) shrinks the cap so
        the byte budget is preserved across model sizes."""
        from omlx.cache.paged_ssd_cache import _compute_max_pending_writes

        cap_small = _compute_max_pending_writes(kv_bytes_per_token=50_000)
        cap_large = _compute_max_pending_writes(kv_bytes_per_token=400_000)

        assert cap_large <= cap_small

    def test_default_args_produce_sensible_cap(self):
        """The default-args path used by static callers must produce
        a cap inside the bounded range."""
        from omlx.cache.paged_ssd_cache import _compute_max_pending_writes

        cap = _compute_max_pending_writes()
        assert 1 <= cap <= 256

    def test_manager_picks_up_per_instance_cap(self, tmp_path):
        """The PagedSSDCacheManager must recompute the cap from its
        constructor args, not just inherit the module-level constant.
        Otherwise a non-default block size silently uses a cap sized
        for the default block size."""
        from omlx.cache.paged_ssd_cache import (
            PagedSSDCacheManager,
            _compute_max_pending_writes,
        )

        mgr_small = PagedSSDCacheManager(
            cache_dir=tmp_path / "small",
            max_size_bytes=1 << 30,
            expected_block_size_tokens=256,
            expected_kv_bytes_per_token=200_000,
        )
        mgr_large = PagedSSDCacheManager(
            cache_dir=tmp_path / "large",
            max_size_bytes=1 << 30,
            expected_block_size_tokens=2048,
            expected_kv_bytes_per_token=200_000,
        )

        # Large blocks must produce a cap no larger than small blocks.
        # Both should equal what the standalone function computes for
        # the same inputs (no surprise rescaling inside __init__).
        assert mgr_small._max_pending_writes == _compute_max_pending_writes(
            block_size_tokens=256, kv_bytes_per_token=200_000
        )
        assert mgr_large._max_pending_writes == _compute_max_pending_writes(
            block_size_tokens=2048, kv_bytes_per_token=200_000
        )
        assert mgr_large._max_pending_writes <= mgr_small._max_pending_writes

        # The write queue must use the instance cap, not the module
        # default — otherwise the formula change is ineffective.
        assert mgr_small._write_queue.maxsize == mgr_small._max_pending_writes
        assert mgr_large._write_queue.maxsize == mgr_large._max_pending_writes

        mgr_small.close()
        mgr_large.close()


class TestSchedulerPlumbsBlockSizeToSSDCache:
    """The Scheduler construction path must plumb its final
    ``paged_cache_block_size`` and a model-derived KV bytes-per-token
    estimate into ``PagedSSDCacheManager.__init__`` — otherwise the
    block-size-aware queue cap silently falls back to defaults sized
    for the wrong workload.

    Regression guard for jundot's review on PR #1627:

      > PagedSSDCacheManager now accepts ``expected_block_size_tokens``
      > / ``expected_kv_bytes_per_token``, but the real scheduler
      > construction path still does not pass them.
    """

    def _make_scheduler(self, tmp_path, block_size_tokens, model_layers):
        """Build a Scheduler with paged SSD cache enabled at the given
        block size and a model whose config exposes ``model_layers``
        layers (so the memory monitor produces a real per-token KV
        estimate rather than its default)."""
        from unittest.mock import MagicMock

        from omlx.scheduler import Scheduler, SchedulerConfig

        # Real numbers so MemoryMonitor.set_model_info accepts them.
        class _Config:
            num_hidden_layers = model_layers
            num_key_value_heads = 8
            num_attention_heads = 32
            head_dim = 192
            hidden_size = 6144

        model = MagicMock()
        model.layers = []
        model.config = _Config()

        tokenizer = MagicMock()
        tokenizer.eos_token_id = 2

        config = SchedulerConfig(
            max_num_seqs=4,
            prefill_step_size=2048,
            paged_cache_block_size=block_size_tokens,
            paged_ssd_cache_dir=str(tmp_path),
            paged_ssd_cache_max_size=1 << 30,
            hot_cache_max_size=0,
            model_name="test-model",
        )
        return Scheduler(model=model, tokenizer=tokenizer, config=config)

    def test_block_size_plumbed_to_ssd_manager(self, tmp_path):
        """Final scheduler block size must reach the SSD cache manager
        — not the 256-token default."""
        model_layers = 32
        sched = self._make_scheduler(
            tmp_path / "blk1024",
            block_size_tokens=1024,
            model_layers=model_layers,
        )

        mgr = sched.paged_ssd_cache_manager

        # The manager stashes the constructor inputs as
        # ``_expected_block_size_tokens`` / ``_expected_kv_bytes_per_token``
        # so the plumbing-regression check is bulletproof regardless
        # of the test host's RAM tier (where the cap math might clamp
        # both default and plumbed inputs to the same floor/ceiling).
        assert mgr._expected_block_size_tokens == 1024, (
            "Scheduler must plumb its final paged_cache_block_size "
            "into PagedSSDCacheManager"
        )

        # The auto-init in ``Scheduler.__init__`` (paired with
        # ``_set_model_info_for_monitor()``) must populate dims from
        # the fixture's model.config so the writer-queue cap formula
        # weighs a real per-token cost. If either step regresses the
        # SSD manager would silently fall back to either the monitor's
        # ~128 KB 7B-class fiction or the manager's own 200 KB default
        # — both miscalibrate the cap on the real model. Assert the
        # construction order didn't regress.
        assert sched.memory_monitor is not None, (
            "Scheduler.__init__ must auto-init MemoryMonitor before "
            "_init_tiered_cache; if this regresses the SSD manager "
            "silently uses its 200 KB default and the cap math is "
            "miscalibrated on every workload"
        )
        assert sched.memory_monitor.has_model_info(), (
            "_set_model_info_for_monitor must populate dims from "
            "model.config BEFORE _init_tiered_cache constructs the "
            "PagedSSDCacheManager; otherwise estimate_block_memory(1) "
            "returns its 7B-class default fiction"
        )

        # The model-derived per-token KV is what should reach the
        # manager. Compute it inline from the fixture's known dims so
        # this assertion can't be satisfied by a tautological pass-
        # through of whatever ``estimate_block_memory`` happens to
        # return — and so it doesn't equal the 200 KB manager default,
        # making a silent fallback to the default visible as a numeric
        # diff in the failure message.
        per_token_bytes = (
            1  # block_size token
            * 8  # num_kv_heads in _make_scheduler fixture
            * 192  # head_dim
            * 2  # dtype_size (float16)
            * 2  # keys + values
        ) * model_layers
        assert mgr._expected_kv_bytes_per_token == per_token_bytes, (
            f"PagedSSDCacheManager received "
            f"expected_kv_bytes_per_token={mgr._expected_kv_bytes_per_token} "
            f"but the model-derived value is {per_token_bytes}; "
            f"a value of 200000 indicates a silent fallback to the "
            f"manager default"
        )
        # And the value the manager received must match
        # ``memory_monitor.estimate_block_memory(1)`` — the documented
        # source. Keep both checks: the inline calc above catches
        # tautological pass-throughs, this one catches drift between
        # the monitor's calc and the scheduler's wiring.
        assert mgr._expected_kv_bytes_per_token == (
            sched.memory_monitor.estimate_block_memory(1)
        ), (
            "Scheduler must plumb memory_monitor.estimate_block_memory"
            "(1) as expected_kv_bytes_per_token"
        )

        # And the cap computed from those plumbed inputs must drive the
        # write queue's maxsize — the cap is only useful if the queue
        # actually enforces it.
        from omlx.cache.paged_ssd_cache import _compute_max_pending_writes

        expected_cap = _compute_max_pending_writes(
            block_size_tokens=mgr._expected_block_size_tokens,
            kv_bytes_per_token=mgr._expected_kv_bytes_per_token,
        )
        assert mgr._max_pending_writes == expected_cap
        assert mgr._write_queue.maxsize == mgr._max_pending_writes

        mgr.close()

    def test_larger_block_size_shrinks_scheduler_cap(self, tmp_path):
        """The same model on the same Mac with a bigger block size
        must produce a smaller pending-writes cap — proves the wiring
        is not stuck at the default for every scheduler instance."""
        sched_small = self._make_scheduler(
            tmp_path / "blk256",
            block_size_tokens=256,
            model_layers=32,
        )
        sched_large = self._make_scheduler(
            tmp_path / "blk2048",
            block_size_tokens=2048,
            model_layers=32,
        )

        cap_small = sched_small.paged_ssd_cache_manager._max_pending_writes
        cap_large = sched_large.paged_ssd_cache_manager._max_pending_writes
        assert cap_large <= cap_small, (
            f"Larger blocks must shrink the cap; got "
            f"cap(2048)={cap_large} > cap(256)={cap_small}"
        )

        sched_small.paged_ssd_cache_manager.close()
        sched_large.paged_ssd_cache_manager.close()

    def test_factory_path_also_plumbs_block_size(self, tmp_path):
        """The cache factory path (used by direct/test callers) must
        also pass ``config.block_size`` through — otherwise direct
        callers silently get the 256-token default regardless of what
        they configured."""
        from omlx.cache.factory import CacheConfig, CacheFactory

        cfg = CacheConfig(
            block_size=1024,
            paged_ssd_cache_dir=tmp_path,
            max_paged_ssd_cache_size=1 << 30,
        )
        mgr = CacheFactory.create_paged_ssd_cache(cfg, model_name="m")
        assert mgr is not None

        from omlx.cache.paged_ssd_cache import _compute_max_pending_writes

        # Factory has no memory monitor, so it leaves
        # ``kv_bytes_per_token`` at the 200 KB default — that's the
        # documented contract for non-Scheduler callers.
        expected = _compute_max_pending_writes(
            block_size_tokens=1024, kv_bytes_per_token=200_000
        )
        assert mgr._max_pending_writes == expected
        mgr.close()


class TestInlineLRUUnlinks:
    """LRU eviction must unlink inline on the calling thread, not enqueue
    ``("unlink", path)`` tasks onto ``_write_queue``.

    The original async-queued design routed eviction unlinks through the
    same bounded queue that carries pending writes. Under sustained save
    pressure, the queue saturated, ``save_block``'s pre-eviction
    ``_write_queue.full()`` short-circuit fired before eviction could
    run, and the cache stayed permanently full once the queue saturated.
    Inlining removes the bounded-queue contention.
    """

    @pytest.fixture
    def mx(self):
        try:
            import mlx.core as mx

            return mx
        except ImportError:
            pytest.skip("MLX not available")

    def _entry_size(self, num_layers=2, seq_len=16, heads=2, head_dim=16):
        # 2 tensors (K+V) per layer, batch=1, float32=4
        return num_layers * 2 * 1 * heads * seq_len * head_dim * 4

    def _save_block(self, mgr, mx, block_hash, num_layers=2):
        cache_data = [
            (mx.zeros((1, 2, 16, 16)), mx.zeros((1, 2, 16, 16)))
            for _ in range(num_layers)
        ]
        return mgr.save_block(
            block_hash=block_hash,
            cache_data=cache_data,
            token_count=16,
            model_name="test-model",
            layer_cache_types=["KVCache"] * num_layers,
        )

    def test_save_waits_on_full_queue_before_inline_fallback(self, tmp_path, mx):
        """A full queue must not short-circuit before tensor extraction.

        The bounded wait gives the writer a chance to drain transient bursts;
        sustained saturation writes the block inline.
        """
        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "full_queue_wait",
            max_size_bytes=1 << 30,
        )
        original_put = mgr._write_queue.put
        original_full = mgr._write_queue.full
        calls: list[float | None] = []

        def fake_put(item, timeout=None, *args, **kwargs):
            calls.append(timeout)
            raise queue.Full

        try:
            mgr._write_queue.full = lambda: True  # type: ignore[method-assign]
            mgr._write_queue.put = fake_put  # type: ignore[method-assign]

            block_hash = b"full_queue_wait"
            assert self._save_block(mgr, mx, block_hash) is True

            assert calls == [1.0]
            stats = mgr.get_stats()
            assert stats.ssd_write_drops == 0
            assert stats.ssd_inline_write_fallbacks == 1
            assert stats.saves == 1
            assert stats.saves_persisted == 1
            assert mgr.load_block(block_hash) is not None
        finally:
            mgr._write_queue.put = original_put  # type: ignore[method-assign]
            mgr._write_queue.full = original_full  # type: ignore[method-assign]
            mgr.close()

    def test_eviction_does_not_enqueue_unlink_tasks(self, tmp_path, mx):
        """Force eviction; assert no ``("unlink", ...)`` items ever enter
        ``_write_queue``. Regression for the original async-queued
        design."""
        entry_size = self._entry_size()
        # Room for ~2 entries; the third save forces eviction of the first.
        max_bytes = entry_size * 2 + 100

        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "inline_eviction",
            max_size_bytes=max_bytes,
        )

        # Sentinel: intercept put_nowait and reject any unlink-shaped tuple.
        original_put_nowait = mgr._write_queue.put_nowait
        unlink_attempts: list = []

        def guard_put_nowait(item):
            if isinstance(item, tuple) and item and item[0] == "unlink":
                unlink_attempts.append(item)
            return original_put_nowait(item)

        mgr._write_queue.put_nowait = guard_put_nowait  # type: ignore[assignment]

        try:
            for i in range(5):
                self._save_block(mgr, mx, f"inline_evict_{i:04d}".encode())

            assert unlink_attempts == [], (
                "Eviction must unlink inline, not enqueue. Found queued "
                f"unlink attempts: {unlink_attempts!r}"
            )
        finally:
            mgr.close()

    def test_eviction_frees_capacity_under_pressure(self, tmp_path, mx):
        """Even when the writer thread is paused (mimicking the
        saturation scenario), eviction must keep the index size within
        the configured cap."""
        entry_size = self._entry_size()
        max_bytes = entry_size * 2 + 100  # holds exactly 2 entries

        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "inline_pressure",
            max_size_bytes=max_bytes,
        )
        try:
            # Save more entries than the cap allows; eviction must keep
            # the index within ``max_bytes``.
            for i in range(8):
                self._save_block(mgr, mx, f"pressure_{i:04d}".encode())

            # Wait briefly for in-flight writes to settle so the index
            # accounting reflects post-eviction state.
            time.sleep(0.05)

            assert mgr._index.total_size <= max_bytes + entry_size, (
                f"Eviction failed to keep total_size ({mgr._index.total_size}) "
                f"near cap ({max_bytes})"
            )
        finally:
            mgr.close()

    def test_inline_eviction_burst_is_capped(self, tmp_path, mx):
        """A large forced eviction is bounded by
        ``_MAX_INLINE_UNLINKS_PER_SAVE``; deferred entries reinsert into
        the index so subsequent saves drain the remainder."""
        from omlx.cache.paged_ssd_cache import _MAX_INLINE_UNLINKS_PER_SAVE

        # Use a large cap initially, then shrink to force a mass-eviction.
        entry_size = self._entry_size()
        n_entries = _MAX_INLINE_UNLINKS_PER_SAVE + 16
        initial_max = entry_size * (n_entries + 2)

        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "inline_burst",
            max_size_bytes=initial_max,
        )
        try:
            for i in range(n_entries):
                self._save_block(mgr, mx, f"burst_{i:04d}".encode())

            # Wait for writes to flush so file_size in the index matches
            # what's on disk.
            time.sleep(0.05)
            count_before = mgr._index.count

            # Shrink the effective cap dramatically. Next eviction must
            # cap its inline burst at _MAX_INLINE_UNLINKS_PER_SAVE.
            mgr._max_size = entry_size  # cap at 1 entry

            # Trigger eviction via a fresh save.
            self._save_block(mgr, mx, b"burst_trigger___")

            # After one save, the index should have shed at most
            # _MAX_INLINE_UNLINKS_PER_SAVE entries. The rest must have
            # been reinserted so subsequent saves can drain them.
            time.sleep(0.05)
            count_after = mgr._index.count
            removed = count_before + 1 - count_after  # +1 for the new save
            assert removed <= _MAX_INLINE_UNLINKS_PER_SAVE, (
                f"Inline burst removed {removed} entries (cap "
                f"{_MAX_INLINE_UNLINKS_PER_SAVE}); ENOSPC-storm protection "
                f"is not in effect"
            )
            assert removed > 0, (
                "No entries were evicted despite the new save crossing "
                "the (shrunken) cap"
            )
        finally:
            mgr.close()

    def test_deferred_eviction_preserves_lru_order(self, tmp_path, monkeypatch):
        """Deferred eviction entries remain older than survivor entries."""
        from omlx.cache import paged_ssd_cache as ssd_cache_module

        monkeypatch.setattr(ssd_cache_module, "_MAX_INLINE_UNLINKS_PER_SAVE", 2)

        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "deferred_lru_order",
            max_size_bytes=1024**2 + 20,
            hot_cache_only=True,
        )
        try:
            for i in range(6):
                block_hash = f"lru_{i}".encode()
                mgr._index.add(
                    PagedSSDBlockMetadata(
                        block_hash=block_hash,
                        file_path=tmp_path / f"{block_hash.hex()}.safetensors",
                        file_size=10,
                        token_count=1,
                        created_at=float(i),
                        last_access=float(i),
                        num_layers=1,
                    )
                )

            mgr._get_effective_max_size = (  # type: ignore[method-assign]
                lambda: 1024**2 + 20
            )
            mgr._enforce_size_limit_for_new_block()

            remaining_lru = [
                metadata.block_hash
                for metadata in mgr._index.get_lru_entries(mgr._index.count)
            ]
            assert remaining_lru == [b"lru_2", b"lru_3", b"lru_4", b"lru_5"]
        finally:
            mgr.close()

    def test_unlink_failure_increments_counter(self, tmp_path, mx):
        """When ``Path.unlink`` raises ``OSError``, the eviction loop
        records the failure in ``evict_unlink_failures`` instead of
        silently dropping the signal."""
        entry_size = self._entry_size()
        max_bytes = entry_size + 100

        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "unlink_fail",
            max_size_bytes=max_bytes,
        )
        try:
            self._save_block(mgr, mx, b"unlink_fail_0001")
            time.sleep(0.05)

            # Patch unlink to raise OSError on the next eviction attempt.
            from pathlib import Path as _Path

            def boom_unlink(self, *args, **kwargs):
                raise OSError("simulated unlink failure")

            with patch.object(_Path, "unlink", boom_unlink):
                # Save a second block; eviction of the first triggers the
                # patched unlink.
                self._save_block(mgr, mx, b"unlink_fail_0002")
                time.sleep(0.05)

            assert mgr._stats["evict_unlink_failures"] >= 1
        finally:
            mgr.close()


class TestSharedHotCacheBudgetClearAllOwners:
    """clear_all_owners reaches managers the budget still pins (orphaned)."""

    class _FakeOwner:
        def __init__(self, budget, n):
            self._budget = budget
            self._n = n
            self.cleared = False

        def clear_hot_cache(self):
            self._budget.forget_owner(self)
            self.cleared = True
            return self._n

    def test_clears_every_tracked_owner(self):
        budget = SharedHotCacheBudget(1 << 20)
        o1 = self._FakeOwner(budget, 3)
        o2 = self._FakeOwner(budget, 4)
        budget.put(o1, b"h1", 100)
        budget.put(o1, b"h2", 100)
        budget.put(o2, b"h3", 200)

        cleared = budget.clear_all_owners()

        assert cleared == 7
        assert o1.cleared and o2.cleared
        assert len(budget._entries) == 0

    def test_empty_budget_is_noop(self):
        assert SharedHotCacheBudget(1 << 20).clear_all_owners() == 0


class TestLayerSignatureSweep:
    """Tests for ``adopt_layer_signature_if_unset`` /
    ``invalidate_stale_layer_signature``.

    The bug: blocks saved before a cache-config change (e.g., TurboQuant
    toggle) carry stale ``layer_cache_types`` metadata, and the
    per-block compatibility check in ``reconstruct_cache`` uses block 0
    as the reference — so a stale block 0 defines the reference forever
    and every newer correctly-typed block trips the mismatch.
    """

    def _make_meta(
        self,
        *,
        block_hash: bytes,
        model_name: str,
        layer_cache_types: list[str] | None,
        num_layers: int = 40,
        block_size: int = 2048,
    ) -> PagedSSDBlockMetadata:
        now = time.time()
        return PagedSSDBlockMetadata(
            block_hash=block_hash,
            file_path=Path("/tmp/never-touched.safetensors"),
            file_size=1024,
            token_count=block_size,
            created_at=now,
            last_access=now,
            num_layers=num_layers,
            model_name=model_name,
            block_size=block_size,
            layer_cache_types=layer_cache_types,
        )

    def _make_manager(
        self,
        tmp_path: Path,
        *,
        model_name: str = "test-model",
        expected_layer_cache_types: list[str] | None = None,
    ) -> PagedSSDCacheManager:
        return PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1 << 30,
            expected_model_name=model_name,
            expected_num_layers=40,
            expected_block_size=2048,
            expected_layer_cache_types=expected_layer_cache_types,
        )

    def test_adopt_sets_signature_first_call(self, tmp_path: Path):
        mgr = self._make_manager(tmp_path)
        assert mgr._expected_layer_cache_types is None

        sig = ["ArraysCache", "ArraysCache", "ArraysCache", "TurboQuantKVCache"]
        adopted = mgr.adopt_layer_signature_if_unset(sig)

        assert adopted is True
        assert mgr._expected_layer_cache_types == sig

    def test_adopt_noop_when_already_set(self, tmp_path: Path):
        original = ["ArraysCache", "KVCache"]
        mgr = self._make_manager(tmp_path, expected_layer_cache_types=original)

        adopted = mgr.adopt_layer_signature_if_unset(
            ["ArraysCache", "TurboQuantKVCache"]
        )

        assert adopted is False
        assert mgr._expected_layer_cache_types == original

    def test_adopt_noop_when_empty(self, tmp_path: Path):
        mgr = self._make_manager(tmp_path)
        assert mgr.adopt_layer_signature_if_unset(None) is False
        assert mgr.adopt_layer_signature_if_unset([]) is False
        assert mgr._expected_layer_cache_types is None

    def test_sweep_drops_blocks_with_stale_signature(self, tmp_path: Path):
        stock = ["ArraysCache", "ArraysCache", "ArraysCache", "KVCache"]
        turbo = [
            "ArraysCache",
            "ArraysCache",
            "ArraysCache",
            "TurboQuantKVCache",
        ]
        mgr = self._make_manager(tmp_path, expected_layer_cache_types=turbo)

        mgr._index.add(
            self._make_meta(
                block_hash=b"01" * 10,
                model_name="test-model",
                layer_cache_types=stock,
                num_layers=4,
            )
        )
        mgr._index.add(
            self._make_meta(
                block_hash=b"02" * 10,
                model_name="test-model",
                layer_cache_types=turbo,
                num_layers=4,
            )
        )

        dropped = mgr.invalidate_stale_layer_signature()

        assert dropped == 1
        assert mgr._index.get(b"01" * 10) is None
        assert mgr._index.get(b"02" * 10) is not None

    def test_sweep_canonicalizes_sized_arrays_cache(self, tmp_path: Path):
        # SizedArraysCache and ArraysCache must compare equal — both are
        # the same on-disk format. The sweep must NOT drop blocks just
        # because one side uses the wrapper class name.
        a = ["ArraysCache", "ArraysCache", "KVCache"]
        b = ["SizedArraysCache", "SizedArraysCache", "KVCache"]
        mgr = self._make_manager(tmp_path, expected_layer_cache_types=a)
        mgr._index.add(
            self._make_meta(
                block_hash=b"aa" * 10,
                model_name="test-model",
                layer_cache_types=b,
                num_layers=3,
            )
        )

        dropped = mgr.invalidate_stale_layer_signature()

        assert dropped == 0
        assert mgr._index.get(b"aa" * 10) is not None

    def test_sweep_canonicalizes_prefill_ready_rotating_cache(self, tmp_path: Path):
        expected = ["KVCache", "RotatingKVCache", "KVCache"]
        stored = ["KVCache", "PrefillReadyRotatingKVCache", "KVCache"]
        mgr = self._make_manager(tmp_path, expected_layer_cache_types=expected)
        mgr._index.add(
            self._make_meta(
                block_hash=b"ac" * 10,
                model_name="test-model",
                layer_cache_types=stored,
                num_layers=3,
            )
        )

        dropped = mgr.invalidate_stale_layer_signature()

        assert dropped == 0
        assert mgr._index.get(b"ac" * 10) is not None

    def test_sweep_drops_legacy_buffered_rotating_cache(self, tmp_path: Path):
        expected = ["KVCache", "RotatingKVCache", "KVCache"]
        stored = ["KVCache", "BufferedRotatingKVCache", "KVCache"]
        mgr = self._make_manager(tmp_path, expected_layer_cache_types=expected)
        mgr._index.add(
            self._make_meta(
                block_hash=b"af" * 10,
                model_name="test-model",
                layer_cache_types=stored,
                num_layers=3,
            )
        )

        dropped = mgr.invalidate_stale_layer_signature()

        assert dropped == 1
        assert mgr._index.get(b"af" * 10) is None

    def test_startup_compat_canonicalizes_sized_arrays_cache(self, tmp_path: Path):
        expected = ["ArraysCache", "ArraysCache", "KVCache"]
        stored = ["SizedArraysCache", "SizedArraysCache", "KVCache"]
        mgr = self._make_manager(tmp_path, expected_layer_cache_types=expected)
        meta = self._make_meta(
            block_hash=b"ab" * 10,
            model_name="test-model",
            layer_cache_types=stored,
            num_layers=40,
        )
        meta.cache_signature = _cache_compat_signature(
            model_name="test-model",
            num_layers=40,
            block_size=2048,
            layer_cache_types=stored,
        )

        assert mgr._is_compatible_block(meta) is True

    def test_startup_compat_canonicalizes_prefill_ready_rotating_cache(
        self, tmp_path: Path
    ):
        expected = ["KVCache", "RotatingKVCache", "KVCache"]
        stored = ["KVCache", "PrefillReadyRotatingKVCache", "KVCache"]
        mgr = self._make_manager(tmp_path, expected_layer_cache_types=expected)
        meta = self._make_meta(
            block_hash=b"ad" * 10,
            model_name="test-model",
            layer_cache_types=stored,
            num_layers=40,
        )
        meta.cache_signature = _cache_compat_signature(
            model_name="test-model",
            num_layers=40,
            block_size=2048,
            layer_cache_types=stored,
        )

        assert mgr._is_compatible_block(meta) is True

    def test_set_expected_layer_signature_replaces_existing(self, tmp_path: Path):
        mgr = self._make_manager(
            tmp_path,
            expected_layer_cache_types=["ArraysCache", "KVCache"],
        )
        mgr._signature_sweep_completed = True

        changed = mgr.set_expected_layer_signature(["ArraysCache", "TurboQuantKVCache"])

        assert changed is True
        assert mgr._expected_layer_cache_types == [
            "ArraysCache",
            "TurboQuantKVCache",
        ]
        assert mgr._signature_sweep_completed is False

    def test_set_expected_layer_signature_noop_for_canonical_match(
        self, tmp_path: Path
    ):
        mgr = self._make_manager(
            tmp_path,
            expected_layer_cache_types=["ArraysCache", "KVCache"],
        )
        mgr._signature_sweep_completed = True

        changed = mgr.set_expected_layer_signature(["SizedArraysCache", "KVCache"])

        assert changed is False
        assert mgr._expected_layer_cache_types == ["SizedArraysCache", "KVCache"]
        assert mgr._signature_sweep_completed is True

    def test_set_expected_layer_signature_noop_for_prefill_ready_rotating_match(
        self, tmp_path: Path
    ):
        mgr = self._make_manager(
            tmp_path,
            expected_layer_cache_types=["KVCache", "RotatingKVCache"],
        )
        mgr._signature_sweep_completed = True

        changed = mgr.set_expected_layer_signature(
            ["KVCache", "PrefillReadyRotatingKVCache"]
        )

        assert changed is False
        assert mgr._expected_layer_cache_types == [
            "KVCache",
            "PrefillReadyRotatingKVCache",
        ]
        assert mgr._signature_sweep_completed is True

    def test_sweep_leaves_other_models_alone(self, tmp_path: Path):
        turbo = ["ArraysCache", "TurboQuantKVCache"]
        stock = ["ArraysCache", "KVCache"]
        mgr = self._make_manager(
            tmp_path, model_name="model-A", expected_layer_cache_types=turbo
        )

        # Other-model block with the "wrong" signature for THIS manager
        # but perfectly valid for model-B.
        mgr._index.add(
            self._make_meta(
                block_hash=b"bb" * 10,
                model_name="model-B",
                layer_cache_types=stock,
                num_layers=2,
            )
        )
        # Same-model stale block.
        mgr._index.add(
            self._make_meta(
                block_hash=b"cc" * 10,
                model_name="model-A",
                layer_cache_types=stock,
                num_layers=2,
            )
        )

        dropped = mgr.invalidate_stale_layer_signature()

        assert dropped == 1
        assert (
            mgr._index.get(b"bb" * 10) is not None
        ), "model-B block must not be touched by model-A's sweep"
        assert mgr._index.get(b"cc" * 10) is None

    def test_sweep_skips_legacy_unnamed_blocks(self, tmp_path: Path):
        turbo = ["ArraysCache", "TurboQuantKVCache"]
        mgr = self._make_manager(
            tmp_path, model_name="model-A", expected_layer_cache_types=turbo
        )
        # No model_name — cannot safely attribute.
        mgr._index.add(
            self._make_meta(
                block_hash=b"dd" * 10,
                model_name="",
                layer_cache_types=["ArraysCache", "KVCache"],
                num_layers=2,
            )
        )

        dropped = mgr.invalidate_stale_layer_signature()

        assert dropped == 0
        assert mgr._index.get(b"dd" * 10) is not None

    def test_sweep_skips_blocks_without_layer_metadata(self, tmp_path: Path):
        turbo = ["ArraysCache", "TurboQuantKVCache"]
        mgr = self._make_manager(
            tmp_path, model_name="model-A", expected_layer_cache_types=turbo
        )
        # Pre-signature block: no layer_cache_types recorded.
        mgr._index.add(
            self._make_meta(
                block_hash=b"ee" * 10,
                model_name="model-A",
                layer_cache_types=None,
                num_layers=2,
            )
        )

        dropped = mgr.invalidate_stale_layer_signature()

        assert dropped == 0
        assert mgr._index.get(b"ee" * 10) is not None

    def test_sweep_idempotent(self, tmp_path: Path):
        turbo = ["ArraysCache", "TurboQuantKVCache"]
        mgr = self._make_manager(
            tmp_path, model_name="model-A", expected_layer_cache_types=turbo
        )
        mgr._index.add(
            self._make_meta(
                block_hash=b"ff" * 10,
                model_name="model-A",
                layer_cache_types=["ArraysCache", "KVCache"],
                num_layers=2,
            )
        )

        assert mgr.invalidate_stale_layer_signature() == 1
        # Subsequent calls do nothing — flag is set.
        assert mgr.invalidate_stale_layer_signature() == 0

    def test_sweep_noop_without_signature(self, tmp_path: Path):
        mgr = self._make_manager(tmp_path)  # signature unset
        mgr._index.add(
            self._make_meta(
                block_hash=b"99" * 10,
                model_name="test-model",
                layer_cache_types=["ArraysCache", "KVCache"],
                num_layers=2,
            )
        )
        assert mgr.invalidate_stale_layer_signature() == 0
        assert mgr._index.get(b"99" * 10) is not None

    def test_sweep_noop_without_model_name(self, tmp_path: Path):
        # Without an expected_model_name we cannot scope safely.
        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1 << 30,
            expected_model_name="",  # explicit unset
            expected_layer_cache_types=["ArraysCache", "KVCache"],
        )
        mgr._index.add(
            self._make_meta(
                block_hash=b"77" * 10,
                model_name="some-model",
                layer_cache_types=["ArraysCache", "TurboQuantKVCache"],
                num_layers=2,
            )
        )
        assert mgr.invalidate_stale_layer_signature() == 0
        assert mgr._index.get(b"77" * 10) is not None

    def test_adopt_resets_sweep_flag(self, tmp_path: Path):
        # If a previous signature got adopted-and-swept, but somehow a new
        # signature gets adopted later (currently the manager treats this
        # as a no-op; we still want the flag to reset if it ever does).
        mgr = self._make_manager(tmp_path)
        mgr.adopt_layer_signature_if_unset(["ArraysCache", "KVCache"])
        mgr._signature_sweep_completed = True
        # Same call signature — already set, returns False, flag stays.
        assert mgr.adopt_layer_signature_if_unset(["ArraysCache", "KVCache"]) is False
        assert mgr._signature_sweep_completed is True


class TestTurboquantBitsSignature:
    """#2045: the cache signature must key the TurboQuant bit depth.

    TurboQuant packed state width is ceil(head_dim * bits / 32), so blocks
    written at different bit depths are shape-incompatible. Without the
    bits field in the signature, 4-bit and 6-bit blocks for the same model
    coexist and crash ``_concat_state_batch`` when their requests are
    batched together.
    """

    TURBO = ["ArraysCache", "TurboQuantKVCache", "TurboQuantKVCache", "KVCache"]

    def _make_meta(
        self,
        *,
        block_hash: bytes,
        layer_cache_types: list[str] | None,
        turboquant_kv_bits: float | None = None,
        cache_signature: str | None = None,
        model_name: str = "test-model",
        num_layers: int = 4,
        block_size: int = 2048,
    ) -> PagedSSDBlockMetadata:
        now = time.time()
        if cache_signature is None:
            cache_signature = _cache_compat_signature(
                model_name=model_name,
                num_layers=num_layers,
                block_size=block_size,
                layer_cache_types=layer_cache_types,
                turboquant_kv_bits=turboquant_kv_bits,
            )
        return PagedSSDBlockMetadata(
            block_hash=block_hash,
            file_path=Path("/tmp/never-touched.safetensors"),
            file_size=1024,
            token_count=block_size,
            created_at=now,
            last_access=now,
            num_layers=num_layers,
            model_name=model_name,
            block_size=block_size,
            layer_cache_types=layer_cache_types,
            cache_signature=cache_signature,
        )

    def _make_manager(
        self,
        tmp_path: Path,
        *,
        layer_cache_types: list[str] | None = None,
        turboquant_kv_bits: float | None = None,
    ) -> PagedSSDCacheManager:
        # Expectations are set through set_expected_layer_signature — the
        # same call the scheduler's refresh makes — so these tests exercise
        # the path production actually runs.
        manager = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=1 << 30,
            expected_model_name="test-model",
            expected_num_layers=4,
            expected_block_size=2048,
        )
        if layer_cache_types is not None:
            manager.set_expected_layer_signature(
                layer_cache_types, turboquant_kv_bits=turboquant_kv_bits
            )
        return manager

    def test_signature_unchanged_when_bits_none(self):
        # Non-TurboQuant signatures must stay byte-identical to the previous
        # format so existing caches of non-TurboQuant models keep validating.
        old_format = json.dumps(
            {
                "model_name": "m",
                "num_layers": 4,
                "block_size": 2048,
                "layer_cache_types": ["KVCache"],
            },
            sort_keys=True,
            separators=(",", ":"),
        )
        assert (
            _cache_compat_signature(
                model_name="m",
                num_layers=4,
                block_size=2048,
                layer_cache_types=["KVCache"],
            )
            == old_format
        )

    def test_signature_includes_bits_when_set(self):
        sig = _cache_compat_signature(
            model_name="m",
            num_layers=4,
            block_size=2048,
            layer_cache_types=self.TURBO,
            turboquant_kv_bits=6,
        )
        assert json.loads(sig)["turboquant_kv_bits"] == 6.0

    def test_signature_bits_helper_tolerates_non_dict_json(self):
        # A corrupted signature that parses as a JSON scalar/list must read
        # as "no recorded depth", not raise — an AttributeError here would
        # abort the stale-signature sweep and leave mixed-width blocks live.
        assert _signature_turboquant_bits("123") is None
        assert _signature_turboquant_bits("[1, 2]") is None
        assert _signature_turboquant_bits("null") is None
        assert _signature_turboquant_bits("{not json") is None
        assert _signature_turboquant_bits("") is None

    def test_block_bits_read_from_meta_state(self):
        # TurboQuant meta_state is (offset, bits, seed, ...) — the same tuple
        # the restore path reads.
        assert (
            _block_turboquant_bits(self.TURBO, [(), (256, 4.0, 0), (256, 4.0, 0), ()])
            == 4.0
        )
        # Batch-form class name counts too.
        assert (
            _block_turboquant_bits(
                ["BatchTurboQuantKVCache"], [(256, 6.0, 0)]
            )
            == 6.0
        )
        # Non-TurboQuant layouts and absent meta yield None.
        assert _block_turboquant_bits(["KVCache"], [(256,)]) is None
        assert _block_turboquant_bits(self.TURBO, None) is None
        assert _block_turboquant_bits(self.TURBO, [(), (), (), ()]) is None

    def test_compat_rejects_other_or_unproven_bit_depth(self, tmp_path: Path):
        mgr = self._make_manager(
            tmp_path, layer_cache_types=self.TURBO, turboquant_kv_bits=6
        )

        matching = self._make_meta(
            block_hash=b"01" * 10, layer_cache_types=self.TURBO, turboquant_kv_bits=6
        )
        other_depth = self._make_meta(
            block_hash=b"02" * 10, layer_cache_types=self.TURBO, turboquant_kv_bits=4
        )
        # 0.4.x-era block: TurboQuant layout but no bits field recorded.
        legacy = self._make_meta(block_hash=b"03" * 10, layer_cache_types=self.TURBO)
        no_signature = self._make_meta(
            block_hash=b"04" * 10, layer_cache_types=self.TURBO, cache_signature=""
        )

        assert mgr._is_compatible_block(matching) is True
        assert mgr._is_compatible_block(other_depth) is False
        assert mgr._is_compatible_block(legacy) is False
        assert mgr._is_compatible_block(no_signature) is False

    def test_compat_unchanged_when_bits_not_expected(self, tmp_path: Path):
        # TurboQuant disabled: legacy and signature-less blocks keep loading
        # exactly as before (zero behavioral change).
        mgr = self._make_manager(tmp_path, layer_cache_types=self.TURBO)

        legacy = self._make_meta(block_hash=b"05" * 10, layer_cache_types=self.TURBO)
        no_signature = self._make_meta(
            block_hash=b"06" * 10, layer_cache_types=self.TURBO, cache_signature=""
        )

        assert mgr._is_compatible_block(legacy) is True
        assert mgr._is_compatible_block(no_signature) is True

    def test_sweep_drops_other_bit_depth(self, tmp_path: Path):
        mgr = self._make_manager(
            tmp_path, layer_cache_types=self.TURBO, turboquant_kv_bits=6
        )
        mgr._index.add(
            self._make_meta(
                block_hash=b"11" * 10,
                layer_cache_types=self.TURBO,
                turboquant_kv_bits=4,
            )
        )
        mgr._index.add(
            self._make_meta(
                block_hash=b"12" * 10,
                layer_cache_types=self.TURBO,
                turboquant_kv_bits=6,
            )
        )
        mgr._index.add(
            self._make_meta(block_hash=b"13" * 10, layer_cache_types=self.TURBO)
        )
        # Same-model block with no layout metadata at all: with a depth
        # expected it can no more prove its width than its layout.
        mgr._index.add(
            self._make_meta(
                block_hash=b"14" * 10, layer_cache_types=None, cache_signature=""
            )
        )

        dropped = mgr.invalidate_stale_layer_signature()

        assert dropped == 3
        assert mgr._index.get(b"11" * 10) is None  # 4-bit under 6-bit expectation
        assert mgr._index.get(b"12" * 10) is not None  # matching depth survives
        assert mgr._index.get(b"13" * 10) is None  # unproven depth
        assert mgr._index.get(b"14" * 10) is None  # no metadata, depth expected

    def test_sweep_keeps_legacy_blocks_when_bits_not_expected(self, tmp_path: Path):
        mgr = self._make_manager(tmp_path, layer_cache_types=self.TURBO)
        mgr._index.add(
            self._make_meta(block_hash=b"15" * 10, layer_cache_types=self.TURBO)
        )
        # No layout metadata, no depth expectation: skip rather than guess
        # (pre-existing behavior preserved).
        mgr._index.add(
            self._make_meta(
                block_hash=b"16" * 10, layer_cache_types=None, cache_signature=""
            )
        )

        assert mgr.invalidate_stale_layer_signature() == 0
        assert mgr._index.get(b"15" * 10) is not None
        assert mgr._index.get(b"16" * 10) is not None

    def test_batch_and_single_turboquant_names_compare_equal(self, tmp_path: Path):
        # The save path records whichever class name it extracted; the
        # predicted layout always says "TurboQuantKVCache". Batch-form
        # blocks persist the same packed per-request state and must not be
        # swept (that would cold-start the cache every restart).
        batch_form = [
            "ArraysCache",
            "BatchTurboQuantKVCache",
            "BatchTurboQuantKVCache",
            "KVCache",
        ]
        mgr = self._make_manager(
            tmp_path, layer_cache_types=self.TURBO, turboquant_kv_bits=6
        )
        mgr._index.add(
            self._make_meta(
                block_hash=b"21" * 10,
                layer_cache_types=batch_form,
                turboquant_kv_bits=6,
            )
        )

        assert mgr.invalidate_stale_layer_signature() == 0
        assert mgr._index.get(b"21" * 10) is not None
        assert (
            mgr._is_compatible_block(
                self._make_meta(
                    block_hash=b"22" * 10,
                    layer_cache_types=batch_form,
                    turboquant_kv_bits=6,
                )
            )
            is True
        )

    def test_bits_change_alone_triggers_sweep(self, tmp_path: Path):
        mgr = self._make_manager(
            tmp_path, layer_cache_types=self.TURBO, turboquant_kv_bits=4
        )
        mgr._signature_sweep_completed = True

        changed = mgr.set_expected_layer_signature(self.TURBO, turboquant_kv_bits=6)

        assert changed is True
        assert mgr._expected_turboquant_kv_bits == 6.0
        assert mgr._signature_sweep_completed is False

    def test_same_bits_and_types_is_noop(self, tmp_path: Path):
        mgr = self._make_manager(
            tmp_path, layer_cache_types=self.TURBO, turboquant_kv_bits=6
        )
        mgr._signature_sweep_completed = True

        changed = mgr.set_expected_layer_signature(self.TURBO, turboquant_kv_bits=6)

        assert changed is False
        assert mgr._signature_sweep_completed is True

    @pytest.mark.skipif(not _has_mlx(), reason="requires MLX")
    def test_save_path_stamps_block_bits_in_signature(self, tmp_path: Path):
        # Drive the REAL save path: the persisted signature must record the
        # depth from the block's own TurboQuant meta_state (observation),
        # not just the manager's expectation.
        import mlx.core as mx

        mgr = self._make_manager(
            tmp_path,
            layer_cache_types=["TurboQuantKVCache", "KVCache"],
            turboquant_kv_bits=6,
        )
        try:
            keys = mx.zeros((1, 2, 8, 4), dtype=mx.float16)
            values = mx.zeros((1, 2, 8, 4), dtype=mx.float16)
            saved = mgr.save_block(
                block_hash=b"77" * 16,
                cache_data=[(keys, values), (keys, values)],
                token_count=8,
                model_name="test-model",
                layer_cache_types=["TurboQuantKVCache", "KVCache"],
                layer_meta_states=[(8, 6.0, 0), (8,)],
            )
            assert saved is True
            # Wait for the background writer to flush the file to disk.
            file_path = mgr._get_file_path(b"77" * 16)
            for _ in range(50):
                if file_path.exists():
                    break
                time.sleep(0.1)
            assert file_path.exists(), "background writer never produced the file"

            files = list((tmp_path / "ssd_cache").rglob("*.safetensors"))
            assert len(files) == 1
            _, metadata = mx.load(str(files[0]), return_metadata=True)
            payload = json.loads(metadata["cache_signature"])
            assert payload["turboquant_kv_bits"] == 6.0
        finally:
            mgr.close()