//! Chunk cache with hash-based index and LRU eviction.
//!
//! The [`ChunkCache`] avoids re-traversing B-trees on repeated reads of chunked
//! datasets.  On first access it scans the B-tree once and builds a
//! `HashMap<ChunkCoord, ChunkInfo>` (the *chunk index*).  Decompressed chunk
//! data is cached with LRU eviction controlled by a byte-budget.

extern crate alloc;

#[cfg(not(feature = "std"))]
use alloc::{vec, vec::Vec};

use core::ops::{Deref, DerefMut};

#[cfg(not(feature = "std"))]
use alloc::collections::BTreeMap;
#[cfg(feature = "std")]
use std::collections::HashMap;
#[cfg(feature = "std")]
use std::sync::Arc;

use crate::chunk_index::{ChunkIndex, ChunkLayout};
use crate::chunked_read::ChunkInfo;

// ---------------------------------------------------------------------------
// Cache-line alignment constants (TVL — Tensor Virtualization Layout)
// ---------------------------------------------------------------------------

/// Cache line size in bytes for the target architecture.
///
/// ARM64 uses 128-byte cache lines; x86_64 uses 64-byte. We align all chunk
/// buffers to this boundary so SIMD operations can assume aligned input.
#[cfg(target_arch = "aarch64")]
pub const CACHE_LINE_SIZE: usize = 128;

#[cfg(target_arch = "x86_64")]
pub const CACHE_LINE_SIZE: usize = 64;

#[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
pub const CACHE_LINE_SIZE: usize = 64;

/// Round `size` up to the next multiple of [`CACHE_LINE_SIZE`].
#[inline]
pub fn align_to_cache_line(size: usize) -> usize {
    (size + CACHE_LINE_SIZE - 1) & !(CACHE_LINE_SIZE - 1)
}

// ---------------------------------------------------------------------------
// CacheAlignedBuffer
// ---------------------------------------------------------------------------

/// A byte buffer whose data pointer is aligned to [`CACHE_LINE_SIZE`].
///
/// This enables SIMD operations to use aligned loads/stores when processing
/// chunk data, avoiding the penalty of misaligned memory accesses.
///
/// The buffer is backed by `core::alloc::Layout`-controlled allocation and
/// uses `alloc::alloc` for the actual allocation, making it compatible with
/// `no_std` (requires the `alloc` crate). It dereferences to `&[u8]` /
/// `&mut [u8]` for seamless use.
pub struct CacheAlignedBuffer {
    ptr: *mut u8,
    len: usize,
    capacity: usize,
}

// SAFETY: The raw pointer is exclusively owned — no aliasing.
unsafe impl Send for CacheAlignedBuffer {}
// SAFETY: `CacheAlignedBuffer` exposes its contents only via `&[u8]`/`&mut
// [u8]` through the ordinary borrow-checked `Deref`/`DerefMut` impls below —
// the same access pattern as `Vec<u8>`, which is `Sync`. Needed so
// `Arc<CacheAlignedBuffer>` (used by the chunk cache) is itself `Send`.
unsafe impl Sync for CacheAlignedBuffer {}

impl CacheAlignedBuffer {
    /// Allocate a new cache-line-aligned buffer of exactly `len` bytes,
    /// initialized to zero.
    pub fn zeroed(len: usize) -> Self {
        if len == 0 {
            return Self {
                ptr: core::ptr::NonNull::dangling().as_ptr(),
                len: 0,
                capacity: 0,
            };
        }
        let capacity = align_to_cache_line(len);
        let layout = core::alloc::Layout::from_size_align(capacity, CACHE_LINE_SIZE)
            .expect("invalid layout");
        // SAFETY: layout has non-zero size.
        let ptr = unsafe { alloc::alloc::alloc_zeroed(layout) };
        if ptr.is_null() {
            alloc::alloc::handle_alloc_error(layout);
        }
        Self { ptr, len, capacity }
    }

    /// Create a cache-line-aligned copy of an existing byte slice.
    pub fn from_slice(data: &[u8]) -> Self {
        let mut buf = Self::zeroed(data.len());
        buf.as_mut_slice()[..data.len()].copy_from_slice(data);
        buf
    }

    /// Create from an existing `Vec<u8>`, copying into an aligned allocation.
    pub fn from_vec(v: Vec<u8>) -> Self {
        Self::from_slice(&v)
    }

    /// The length of the valid data (may be less than capacity).
    #[inline]
    pub fn len(&self) -> usize {
        self.len
    }

    /// Whether the buffer is empty.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.len == 0
    }

    /// The underlying aligned pointer.
    #[inline]
    pub fn as_ptr(&self) -> *const u8 {
        self.ptr
    }

    /// Mutable pointer to the data.
    #[inline]
    pub fn as_mut_ptr(&mut self) -> *mut u8 {
        self.ptr
    }

    /// Borrow as a byte slice.
    #[inline]
    pub fn as_slice(&self) -> &[u8] {
        if self.len == 0 {
            return &[];
        }
        // SAFETY: ptr is valid for `len` bytes and properly aligned.
        unsafe { core::slice::from_raw_parts(self.ptr, self.len) }
    }

    /// Borrow as a mutable byte slice.
    #[inline]
    pub fn as_mut_slice(&mut self) -> &mut [u8] {
        if self.len == 0 {
            return &mut [];
        }
        // SAFETY: ptr is valid for `len` bytes and properly aligned.
        unsafe { core::slice::from_raw_parts_mut(self.ptr, self.len) }
    }

    /// Convert to a `Vec<u8>` (copies data into a standard allocation).
    pub fn to_vec(&self) -> Vec<u8> {
        self.as_slice().to_vec()
    }

    /// Returns `true` if the data pointer is aligned to `CACHE_LINE_SIZE`.
    #[inline]
    pub fn is_aligned(&self) -> bool {
        self.len == 0 || (self.ptr as usize).is_multiple_of(CACHE_LINE_SIZE)
    }
}

impl Drop for CacheAlignedBuffer {
    fn drop(&mut self) {
        if self.capacity > 0 {
            let layout = core::alloc::Layout::from_size_align(self.capacity, CACHE_LINE_SIZE)
                .expect("invalid layout");
            // SAFETY: ptr was allocated with this layout.
            unsafe { alloc::alloc::dealloc(self.ptr, layout) };
        }
    }
}

impl Clone for CacheAlignedBuffer {
    fn clone(&self) -> Self {
        Self::from_slice(self.as_slice())
    }
}

impl Deref for CacheAlignedBuffer {
    type Target = [u8];
    #[inline]
    fn deref(&self) -> &[u8] {
        self.as_slice()
    }
}

impl DerefMut for CacheAlignedBuffer {
    #[inline]
    fn deref_mut(&mut self) -> &mut [u8] {
        self.as_mut_slice()
    }
}

impl core::fmt::Debug for CacheAlignedBuffer {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("CacheAlignedBuffer")
            .field("len", &self.len)
            .field("capacity", &self.capacity)
            .field("aligned", &self.is_aligned())
            .finish()
    }
}

/// Coordinate key for a chunk -- the N-dimensional offset vector.
pub type ChunkCoord = Vec<u64>;

// ---------------------------------------------------------------------------
// ChunkCache and supporting types require std::sync::Mutex.
// ---------------------------------------------------------------------------

/// Default maximum bytes of decompressed chunk data to cache (16 MiB).
///
/// Increased from the HDF5 C library default of 1 MiB to better accommodate
/// modern workloads with large embedding datasets and high-dimensional chunks.
pub const DEFAULT_CACHE_BYTES: usize = 16 * 1024 * 1024; // 16 MiB

/// Default maximum number of cached decompressed chunks.
///
/// 521 is prime, which provides better hash distribution for the chunk
/// coordinate map and reduces collision chains compared to power-of-two sizes.
pub const DEFAULT_MAX_SLOTS: usize = 521;

/// Most datasets whose chunk index a [`ChunkCache`] keeps at once.
pub const MAX_INDEXED_DATASETS: usize = 64;

/// Most chunk-index entries, summed over all datasets, a [`ChunkCache`] keeps.
/// Least-recently-used datasets' indexes are dropped past this (the dataset
/// being read is always kept), so a file with many or huge chunked datasets
/// cannot grow the cache without bound.
pub const MAX_INDEXED_CHUNKS: usize = 1 << 20;

/// The dataset key the address-less (legacy) methods use when
/// [`ChunkCache::ensure_dataset`] has not been called.
#[cfg(feature = "std")]
const UNBOUND_DATASET: u64 = u64::MAX;

// ---------------------------------------------------------------------------
// LRU entry
// ---------------------------------------------------------------------------

/// Decompressed chunks are keyed by dataset *and* coordinate: every chunked
/// dataset has a chunk at (0, 0, ...), so the coordinate alone is ambiguous.
#[cfg(feature = "std")]
type SlotKey = (u64, ChunkCoord);

#[cfg(feature = "std")]
struct CachedChunk {
    key: SlotKey,
    /// Shared so a cache hit is a refcount bump, not a copy of the whole
    /// (potentially large) decompressed chunk.
    data: Arc<CacheAlignedBuffer>,
    /// Monotonically increasing access counter for LRU ordering.
    last_access: u64,
}

/// Per-dataset index state.
#[cfg(feature = "std")]
#[derive(Default)]
struct DatasetEntry {
    /// Chunk coordinate -> ChunkInfo (offset + size in file).
    index: Option<Arc<HashMap<ChunkCoord, ChunkInfo>>>,
    /// Pre-built chunk index for O(1) coordinate lookups.
    chunk_index: Option<Arc<ChunkIndex>>,
    /// Pre-computed chunk layout for fast assembly.
    chunk_layout: Option<Arc<ChunkLayout>>,
    /// Tick of the last use, for dropping the least recently used dataset.
    last_used: u64,
}

#[cfg(feature = "std")]
impl DatasetEntry {
    fn weight(&self) -> usize {
        self.index.as_ref().map_or(0, |m| m.len())
            + self.chunk_index.as_ref().map_or(0, |c| c.num_chunks())
    }
}

// ---------------------------------------------------------------------------
// ChunkCache
// ---------------------------------------------------------------------------

/// A per-file chunk cache: chunk indexes per dataset, plus an LRU of
/// decompressed chunks, all keyed by dataset.
///
/// A dataset is identified by the address of its chunk index (B-tree, fixed
/// or extensible array, ...), which is unique within a file. Every method
/// that takes an `addr` works on that dataset only, so threads reading
/// different datasets through one shared cache never see each other's
/// chunks. The address-less methods (`has_index`, `populate_index`,
/// `get_decompressed`, ...) act on the dataset last bound with
/// [`Self::ensure_dataset`]; that binding is shared state, so concurrent
/// readers must use the `*_in` / `*_for` methods instead (the chunked
/// readers in [`crate::chunked_read`] do).
///
/// Memory is bounded: decompressed data by `max_bytes`/`max_slots` across
/// all datasets, indexes by [`MAX_INDEXED_DATASETS`] and
/// [`MAX_INDEXED_CHUNKS`].
///
/// Only available with the `std` feature because it requires `std::sync::Mutex`.
#[cfg(feature = "std")]
pub struct ChunkCache {
    inner: std::sync::Mutex<CacheInner>,
}

#[cfg(feature = "std")]
struct CacheInner {
    /// Per-dataset chunk indexes, keyed by chunk-index address.
    datasets: HashMap<u64, DatasetEntry>,

    /// Dataset the address-less methods act on (see `ensure_dataset`).
    current: Option<u64>,

    /// LRU cache of decompressed chunk data.
    slots: Vec<CachedChunk>,

    /// Key -> index into `slots`, for O(1) lookup instead of a linear
    /// scan. Kept in sync with `slots` on every insert/evict/clear — in
    /// particular, `slots.swap_remove(i)` moves the last element into slot
    /// `i`, so the moved element's index entry must be updated too.
    slot_index: HashMap<SlotKey, usize>,

    /// Current total bytes of cached decompressed data.
    current_bytes: usize,

    /// Maximum bytes of decompressed data to cache.
    max_bytes: usize,

    /// Maximum number of slots.
    max_slots: usize,

    /// Monotonic counter for LRU ordering.
    tick: u64,

    /// Last accessed chunk (for sequential detection).
    last_coord: Option<SlotKey>,

    /// Access pattern statistics.
    stats: AccessStats,
}

#[cfg(feature = "std")]
impl CacheInner {
    fn current(&self) -> u64 {
        self.current.unwrap_or(UNBOUND_DATASET)
    }

    fn touch(&mut self, addr: u64) -> &mut DatasetEntry {
        self.tick += 1;
        let tick = self.tick;
        let entry = self.datasets.entry(addr).or_default();
        entry.last_used = tick;
        entry
    }

    fn entry(&self, addr: u64) -> Option<&DatasetEntry> {
        self.datasets.get(&addr)
    }

    /// Drop least-recently-used datasets' indexes (never `keep`'s) until the
    /// dataset and chunk-entry budgets hold.
    fn trim_datasets(&mut self, keep: u64) {
        loop {
            let total: usize = self.datasets.values().map(DatasetEntry::weight).sum();
            if self.datasets.len() <= MAX_INDEXED_DATASETS && total <= MAX_INDEXED_CHUNKS {
                return;
            }
            let victim = self
                .datasets
                .iter()
                .filter(|(a, _)| **a != keep)
                .min_by_key(|(_, e)| e.last_used)
                .map(|(a, _)| *a);
            match victim {
                Some(a) => {
                    self.datasets.remove(&a);
                }
                None => return,
            }
        }
    }

    fn get_decompressed(&mut self, addr: u64, coord: &[u64]) -> Option<Arc<CacheAlignedBuffer>> {
        self.tick += 1;
        let tick = self.tick;

        // Track sequential vs random access
        let is_sequential = self.last_coord.as_ref().is_some_and(|(prev_addr, prev)| {
            // Sequential if exactly one dimension changed
            let changes: usize = prev
                .iter()
                .zip(coord.iter())
                .filter(|(a, b)| a != b)
                .count();
            *prev_addr == addr && changes <= 1
        });
        if is_sequential {
            self.stats.sequential_count += 1;
        } else if self.last_coord.is_some() {
            self.stats.random_count += 1;
        }
        let key: SlotKey = (addr, coord.to_vec());
        let found = if let Some(&idx) = self.slot_index.get(&key) {
            self.slots[idx].last_access = tick;
            Some(Arc::clone(&self.slots[idx].data))
        } else {
            None
        };
        self.last_coord = Some(key);
        if let Some(ref data) = found {
            self.stats.hits += 1;
            self.stats.bytes_read += data.len() as u64;
        } else {
            self.stats.misses += 1;
        }
        found
    }

    fn put_decompressed(
        &mut self,
        key: SlotKey,
        data: Arc<CacheAlignedBuffer>,
    ) -> Arc<CacheAlignedBuffer> {
        let data_len = data.len();

        // Don't cache if single chunk exceeds budget — still return the data
        // to the caller, just don't retain it.
        if data_len > self.max_bytes {
            return data;
        }

        // Check if already present
        self.tick += 1;
        let tick = self.tick;
        if let Some(&idx) = self.slot_index.get(&key) {
            self.slots[idx].last_access = tick;
            return Arc::clone(&self.slots[idx].data); // already cached
        }

        // Evict until we have room
        while self.slots.len() >= self.max_slots
            || (self.current_bytes + data_len > self.max_bytes && !self.slots.is_empty())
        {
            // Find LRU slot
            let lru_idx = self
                .slots
                .iter()
                .enumerate()
                .min_by_key(|(_, s)| s.last_access)
                .map(|(i, _)| i)
                .unwrap();
            let removed = self.slots.swap_remove(lru_idx);
            self.slot_index.remove(&removed.key);
            // swap_remove moved the former last element into `lru_idx` (unless
            // it *was* the last element) — fix up that element's index entry.
            if lru_idx < self.slots.len() {
                let moved_key = self.slots[lru_idx].key.clone();
                self.slot_index.insert(moved_key, lru_idx);
            }
            self.current_bytes -= removed.data.len();
            self.stats.evictions += 1;
        }

        self.current_bytes += data_len;
        let new_idx = self.slots.len();
        self.slot_index.insert(key.clone(), new_idx);
        self.slots.push(CachedChunk {
            key,
            data: Arc::clone(&data),
            last_access: tick,
        });
        data
    }
}

/// Access pattern statistics tracked by the chunk cache.
///
/// Updated on each `get_decompressed` / `put_decompressed` call to help
/// the sweep detector understand the workload.
#[cfg(feature = "std")]
#[derive(Debug, Clone, Default)]
pub struct AccessStats {
    /// Number of accesses that followed a sequential pattern.
    pub sequential_count: u64,
    /// Number of accesses that appeared random (non-sequential).
    pub random_count: u64,
    /// Last detected sweep direction description (informational).
    pub sweep_direction: Option<&'static str>,
    /// Number of cache hits (decompressed data found in cache).
    pub hits: u64,
    /// Number of cache misses (decompressed data not in cache).
    pub misses: u64,
    /// Number of LRU evictions performed.
    pub evictions: u64,
    /// Total bytes read from cached data.
    pub bytes_read: u64,
}

#[cfg(feature = "std")]
impl AccessStats {
    /// Cache hit rate as a fraction in [0.0, 1.0].
    ///
    /// Returns 0.0 if no accesses have been recorded.
    pub fn hit_rate(&self) -> f64 {
        let total = self.hits + self.misses;
        if total == 0 {
            0.0
        } else {
            self.hits as f64 / total as f64
        }
    }
}

#[cfg(feature = "std")]
impl ChunkCache {
    /// Create a new chunk cache with default limits (16 MiB, 521 slots).
    pub fn new() -> Self {
        Self::with_capacity(DEFAULT_CACHE_BYTES, DEFAULT_MAX_SLOTS)
    }

    /// Create a new chunk cache with custom byte budget and slot count.
    pub fn with_capacity(max_bytes: usize, max_slots: usize) -> Self {
        Self {
            inner: std::sync::Mutex::new(CacheInner {
                datasets: HashMap::new(),
                current: None,
                slots: Vec::with_capacity(max_slots.min(64)),
                slot_index: HashMap::with_capacity(max_slots.min(64)),
                current_bytes: 0,
                max_bytes,
                max_slots,
                tick: 0,
                last_coord: None,
                stats: AccessStats::default(),
            }),
        }
    }

    fn lock(&self) -> std::sync::MutexGuard<'_, CacheInner> {
        self.inner.lock().unwrap_or_else(|e| e.into_inner())
    }

    /// The most decompressed bytes this cache will hold.
    pub fn max_bytes(&self) -> usize {
        self.lock().max_bytes
    }

    // ----- Dataset-keyed operations (safe to use concurrently) -----

    /// The chunk list of the dataset whose chunk index is at `addr`.
    ///
    /// On the first call for a dataset, `build` scans its chunk index; the
    /// result is kept (offsets truncated to `rank` for the lookup key), so
    /// later calls skip the scan. `build` runs without the cache lock held;
    /// if two threads race to build the same dataset's index, the first
    /// stored one wins and both return equivalent lists.
    pub fn chunks_for<E>(
        &self,
        addr: u64,
        rank: usize,
        build: impl FnOnce() -> Result<Vec<ChunkInfo>, E>,
    ) -> Result<Vec<ChunkInfo>, E> {
        Ok(self
            .index_for(addr, rank, build)?
            .values()
            .cloned()
            .collect())
    }

    fn index_for<E>(
        &self,
        addr: u64,
        rank: usize,
        build: impl FnOnce() -> Result<Vec<ChunkInfo>, E>,
    ) -> Result<Arc<HashMap<ChunkCoord, ChunkInfo>>, E> {
        if let Some(index) = self.lock().touch(addr).index.clone() {
            return Ok(index);
        }
        let chunks = build()?;
        let map: HashMap<ChunkCoord, ChunkInfo> = chunks
            .into_iter()
            .map(|ci| (ci.offsets.iter().take(rank).copied().collect(), ci))
            .collect();
        let mut inner = self.lock();
        let entry = inner.touch(addr);
        let index = Arc::clone(entry.index.get_or_insert_with(|| Arc::new(map)));
        inner.trim_datasets(addr);
        Ok(index)
    }

    /// The pre-computed assembly layout of the dataset at `addr`, building
    /// its chunk index (via `build`, as in [`Self::chunks_for`]) and layout on
    /// first use.
    pub fn chunk_layout_for<E>(
        &self,
        addr: u64,
        rank: usize,
        build: impl FnOnce() -> Result<Vec<ChunkInfo>, E>,
        ds_dims: &[usize],
        chunk_dims: &[usize],
        elem_size: usize,
    ) -> Result<Arc<ChunkLayout>, E> {
        let (layout, chunk_index) = {
            let mut inner = self.lock();
            let entry = inner.touch(addr);
            (entry.chunk_layout.clone(), entry.chunk_index.clone())
        };
        if let Some(layout) = layout {
            return Ok(layout);
        }
        let chunk_index = match chunk_index {
            Some(ci) => ci,
            None => {
                let index = self.index_for(addr, rank, build)?;
                let chunks: Vec<ChunkInfo> = index.values().cloned().collect();
                Arc::new(ChunkIndex::build(&chunks, rank))
            }
        };
        let layout = ChunkLayout::build(&chunk_index, ds_dims, chunk_dims, elem_size);
        let mut inner = self.lock();
        let entry = inner.touch(addr);
        entry.chunk_index.get_or_insert(chunk_index);
        let layout = Arc::clone(entry.chunk_layout.get_or_insert_with(|| Arc::new(layout)));
        inner.trim_datasets(addr);
        Ok(layout)
    }

    /// Cached decompressed chunk at `coord` of the dataset at `addr`.
    ///
    /// O(1) lookup; the clone is an `Arc` refcount bump, not a copy of the
    /// underlying decompressed data.
    pub fn get_decompressed_in(&self, addr: u64, coord: &[u64]) -> Option<Arc<CacheAlignedBuffer>> {
        self.lock().get_decompressed(addr, coord)
    }

    /// Cache decompressed chunk data for `coord` of the dataset at `addr`.
    /// Returns the `Arc`-shared buffer now cached (or already cached).
    pub fn put_decompressed_in(
        &self,
        addr: u64,
        coord: ChunkCoord,
        data: Vec<u8>,
    ) -> Arc<CacheAlignedBuffer> {
        self.put_decompressed_aligned_in(addr, coord, CacheAlignedBuffer::from_vec(data))
    }

    /// [`Self::put_decompressed_in`] for an already-aligned buffer.
    pub fn put_decompressed_aligned_in(
        &self,
        addr: u64,
        coord: ChunkCoord,
        data: CacheAlignedBuffer,
    ) -> Arc<CacheAlignedBuffer> {
        let data = Arc::new(data);
        self.lock().put_decompressed((addr, coord), data)
    }

    /// Record that the given chunk coordinates of the dataset at `addr` are
    /// predicted to be accessed soon (bookkeeping only).
    ///
    /// This does **not** prefetch or pre-decompress anything — it only
    /// checks whether each coordinate is already in the chunk index and
    /// updates access-pattern stats accordingly.
    pub fn prefetch_hint_in(&self, addr: u64, next_coords: &[ChunkCoord]) {
        let mut inner = self.lock();
        let Some(index) = inner.entry(addr).and_then(|e| e.index.clone()) else {
            return;
        };
        let known = next_coords
            .iter()
            .filter(|c| index.contains_key(*c))
            .count();
        inner.stats.sequential_count += known as u64;
    }

    // ----- Address-less operations on the bound dataset -----

    /// Bind the address-less methods to the dataset at chunk-index address
    /// `addr`. Returns `true` if this changed the bound dataset.
    ///
    /// Each dataset's state is kept separately, so switching loses nothing
    /// and never exposes one dataset's index or chunks to another. The
    /// binding itself is shared, though: concurrent readers should use the
    /// `addr`-taking methods rather than bind and then call these.
    pub fn ensure_dataset(&self, addr: u64) -> bool {
        let mut inner = self.lock();
        let changed = inner.current != Some(addr);
        inner.current = Some(addr);
        changed
    }

    /// Returns `true` if the bound dataset's chunk index has been built.
    pub fn has_index(&self) -> bool {
        let inner = self.lock();
        inner
            .entry(inner.current())
            .is_some_and(|e| e.index.is_some())
    }

    /// Build the bound dataset's chunk index from a pre-collected list of
    /// `ChunkInfo`.
    ///
    /// The `rank` parameter is used to truncate offsets to spatial dims only
    /// (B-tree v1 stores rank+1 offsets).
    pub fn populate_index(&self, chunks: &[ChunkInfo], rank: usize) {
        let addr = self.lock().current();
        let _ = self.index_for::<core::convert::Infallible>(addr, rank, || Ok(chunks.to_vec()));
    }

    /// Look up a chunk by its spatial coordinate in the bound dataset's index.
    pub fn lookup_index(&self, coord: &[u64]) -> Option<ChunkInfo> {
        let inner = self.lock();
        inner
            .entry(inner.current())?
            .index
            .as_ref()?
            .get(coord)
            .cloned()
    }

    /// Return all of the bound dataset's indexed chunks (order unspecified).
    pub fn all_indexed_chunks(&self) -> Option<Vec<ChunkInfo>> {
        let inner = self.lock();
        let index = inner.entry(inner.current())?.index.as_ref()?;
        Some(index.values().cloned().collect())
    }

    /// Returns `true` if the bound dataset's `ChunkIndex` has been built.
    pub fn has_chunk_index(&self) -> bool {
        let inner = self.lock();
        inner
            .entry(inner.current())
            .is_some_and(|e| e.chunk_index.is_some())
    }

    /// Build and store the bound dataset's `ChunkIndex`.
    pub fn populate_chunk_index(&self, chunks: &[ChunkInfo], rank: usize) {
        let built = Arc::new(ChunkIndex::build(chunks, rank));
        let mut inner = self.lock();
        let addr = inner.current();
        inner.touch(addr).chunk_index.get_or_insert(built);
        inner.trim_datasets(addr);
    }

    /// Returns `true` if the bound dataset's chunk layout has been computed.
    pub fn has_chunk_layout(&self) -> bool {
        let inner = self.lock();
        inner
            .entry(inner.current())
            .is_some_and(|e| e.chunk_layout.is_some())
    }

    /// Build and store the bound dataset's chunk layout (needs its
    /// `ChunkIndex`; does nothing without one).
    pub fn populate_chunk_layout(&self, ds_dims: &[usize], chunk_dims: &[usize], elem_size: usize) {
        let mut inner = self.lock();
        let addr = inner.current();
        let entry = inner.touch(addr);
        if entry.chunk_layout.is_some() {
            return;
        }
        if let Some(idx) = entry.chunk_index.clone() {
            entry.chunk_layout = Some(Arc::new(ChunkLayout::build(
                &idx, ds_dims, chunk_dims, elem_size,
            )));
        }
    }

    /// Execute a function with a reference to the bound dataset's chunk
    /// layout. Returns `None` if the layout hasn't been computed yet.
    pub fn with_chunk_layout<F, R>(&self, f: F) -> Option<R>
    where
        F: FnOnce(&ChunkLayout) -> R,
    {
        let layout = {
            let inner = self.lock();
            inner.entry(inner.current())?.chunk_layout.clone()?
        };
        Some(f(&layout))
    }

    /// Try to get cached decompressed data for a chunk of the bound dataset.
    ///
    /// Returns an owned copy; prefer [`Self::get_decompressed_aligned`] when
    /// an `Arc`-shared buffer works for the caller.
    pub fn get_decompressed(&self, coord: &[u64]) -> Option<Vec<u8>> {
        self.get_decompressed_aligned(coord)
            .map(|arc| arc.as_slice().to_vec())
    }

    /// Reference-counted cached buffer for a chunk of the bound dataset.
    pub fn get_decompressed_aligned(&self, coord: &[u64]) -> Option<Arc<CacheAlignedBuffer>> {
        let mut inner = self.lock();
        let addr = inner.current();
        inner.get_decompressed(addr, coord)
    }

    /// Insert decompressed chunk data for the bound dataset into the LRU
    /// cache, returning the `Arc`-shared buffer now cached.
    pub fn put_decompressed(&self, coord: ChunkCoord, data: Vec<u8>) -> Arc<CacheAlignedBuffer> {
        self.put_decompressed_aligned(coord, CacheAlignedBuffer::from_vec(data))
    }

    /// Insert an already-aligned buffer for the bound dataset.
    pub fn put_decompressed_aligned(
        &self,
        coord: ChunkCoord,
        data: CacheAlignedBuffer,
    ) -> Arc<CacheAlignedBuffer> {
        let data = Arc::new(data);
        let mut inner = self.lock();
        let addr = inner.current();
        inner.put_decompressed((addr, coord), data)
    }

    /// [`Self::prefetch_hint_in`] for the bound dataset.
    pub fn prefetch_hint(&self, next_coords: &[ChunkCoord]) {
        let addr = self.lock().current();
        self.prefetch_hint_in(addr, next_coords);
    }

    // ----- Whole-cache operations -----

    /// Clear the entire cache (indexes + decompressed data + stats).
    pub fn clear(&self) {
        let mut inner = self.lock();
        inner.datasets.clear();
        inner.current = None;
        inner.slots.clear();
        inner.slot_index.clear();
        inner.current_bytes = 0;
        inner.tick = 0;
        inner.last_coord = None;
        inner.stats = AccessStats::default();
    }

    /// Return the current access pattern statistics.
    pub fn access_stats(&self) -> AccessStats {
        self.lock().stats.clone()
    }

    /// Update the sweep direction label in the access stats.
    pub fn set_sweep_direction(&self, direction: &'static str) {
        self.lock().stats.sweep_direction = Some(direction);
    }

    /// Number of decompressed chunks currently cached (all datasets).
    pub fn cached_chunk_count(&self) -> usize {
        self.lock().slots.len()
    }

    /// Total bytes of decompressed data currently cached (all datasets).
    pub fn cached_bytes(&self) -> usize {
        self.lock().current_bytes
    }

    /// Number of datasets whose chunk index is currently kept.
    pub fn indexed_dataset_count(&self) -> usize {
        self.lock().datasets.len()
    }
}

#[cfg(feature = "std")]
impl Default for ChunkCache {
    fn default() -> Self {
        Self::new()
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;

    fn make_chunk(offsets: Vec<u64>, address: u64, size: u32) -> ChunkInfo {
        ChunkInfo {
            chunk_size: size,
            filter_mask: 0,
            offsets,
            address,
        }
    }

    #[test]
    fn index_populate_and_lookup() {
        let cache = ChunkCache::new();
        let chunks = vec![
            make_chunk(vec![0, 0, 0], 0x1000, 80),
            make_chunk(vec![10, 0, 0], 0x2000, 80),
        ];
        cache.populate_index(&chunks, 2); // rank=2, truncate to [0,0] and [10,0]
        assert!(cache.has_index());

        let c0 = cache.lookup_index(&[0, 0]).unwrap();
        assert_eq!(c0.address, 0x1000);

        let c1 = cache.lookup_index(&[10, 0]).unwrap();
        assert_eq!(c1.address, 0x2000);

        assert!(cache.lookup_index(&[5, 0]).is_none());
    }

    #[test]
    fn decompressed_cache_hit() {
        let cache = ChunkCache::new();
        cache.put_decompressed(vec![0, 0], vec![1, 2, 3, 4]);
        let got = cache.get_decompressed(&[0, 0]).unwrap();
        assert_eq!(got, vec![1, 2, 3, 4]);
    }

    #[test]
    fn lru_eviction_by_slots() {
        let cache = ChunkCache::with_capacity(1024 * 1024, 2); // max 2 slots

        cache.put_decompressed(vec![0], vec![1; 10]);
        cache.put_decompressed(vec![1], vec![2; 10]);
        assert_eq!(cache.cached_chunk_count(), 2);

        // Access slot 0 to make it more recent
        cache.get_decompressed(&[0]);

        // Insert slot 2 — should evict slot 1 (LRU)
        cache.put_decompressed(vec![2], vec![3; 10]);
        assert_eq!(cache.cached_chunk_count(), 2);

        assert!(cache.get_decompressed(&[0]).is_some());
        assert!(cache.get_decompressed(&[1]).is_none()); // evicted
        assert!(cache.get_decompressed(&[2]).is_some());
    }

    #[test]
    fn lru_eviction_by_bytes() {
        let cache = ChunkCache::with_capacity(50, 100); // 50 bytes max

        cache.put_decompressed(vec![0], vec![0; 20]);
        cache.put_decompressed(vec![1], vec![0; 20]);
        assert_eq!(cache.cached_bytes(), 40);

        // This needs 20 bytes but only 10 free — evict LRU
        cache.put_decompressed(vec![2], vec![0; 20]);
        assert!(cache.cached_bytes() <= 50);
        assert!(cache.get_decompressed(&[0]).is_none()); // evicted (LRU)
    }

    #[test]
    fn oversized_chunk_not_cached() {
        let cache = ChunkCache::with_capacity(10, 16);
        cache.put_decompressed(vec![0], vec![0; 100]); // too big
        assert_eq!(cache.cached_chunk_count(), 0);
    }

    #[test]
    fn clear_resets_everything() {
        let cache = ChunkCache::new();
        let chunks = vec![make_chunk(vec![0, 0], 0x1000, 80)];
        cache.populate_index(&chunks, 1);
        cache.put_decompressed(vec![0], vec![1, 2, 3]);

        cache.clear();
        assert!(!cache.has_index());
        assert_eq!(cache.cached_chunk_count(), 0);
        assert_eq!(cache.cached_bytes(), 0);
    }

    #[test]
    fn datasets_sharing_coordinates_stay_separate() {
        let cache = ChunkCache::new();
        let a = vec![make_chunk(vec![0, 0], 0x100, 8)];
        let b = vec![make_chunk(vec![0, 0], 0x900, 8)];
        let got_a = cache.chunks_for::<()>(1, 1, || Ok(a.clone())).unwrap();
        let got_b = cache.chunks_for::<()>(2, 1, || Ok(b.clone())).unwrap();
        assert_eq!(got_a[0].address, 0x100);
        assert_eq!(got_b[0].address, 0x900);
        // Built once per dataset: a second lookup doesn't call the builder.
        let again = cache
            .chunks_for::<()>(1, 1, || panic!("index rebuilt"))
            .unwrap();
        assert_eq!(again[0].address, 0x100);

        cache.put_decompressed_in(1, vec![0], vec![1; 4]);
        cache.put_decompressed_in(2, vec![0], vec![2; 4]);
        assert_eq!(
            cache.get_decompressed_in(1, &[0]).unwrap().as_slice(),
            &[1; 4]
        );
        assert_eq!(
            cache.get_decompressed_in(2, &[0]).unwrap().as_slice(),
            &[2; 4]
        );
        assert!(cache.get_decompressed_in(3, &[0]).is_none());
        assert_eq!(cache.cached_chunk_count(), 2);

        // The bound-dataset methods see only the bound dataset.
        cache.ensure_dataset(2);
        assert_eq!(cache.lookup_index(&[0]).unwrap().address, 0x900);
        assert_eq!(cache.get_decompressed(&[0]).unwrap(), vec![2; 4]);
    }

    #[test]
    fn dataset_indexes_are_bounded() {
        let cache = ChunkCache::new();
        for addr in 0..(MAX_INDEXED_DATASETS as u64 + 10) {
            cache
                .chunks_for::<()>(addr, 1, || Ok(vec![make_chunk(vec![0], addr, 8)]))
                .unwrap();
        }
        assert_eq!(cache.indexed_dataset_count(), MAX_INDEXED_DATASETS);

        // One huge index evicts the others but is itself kept.
        let huge: Vec<ChunkInfo> = (0..MAX_INDEXED_CHUNKS as u64)
            .map(|i| make_chunk(vec![i], i, 8))
            .collect();
        let got = cache.chunks_for::<()>(9999, 1, || Ok(huge)).unwrap();
        assert_eq!(got.len(), MAX_INDEXED_CHUNKS);
        assert_eq!(cache.indexed_dataset_count(), 1);
    }

    #[test]
    fn concurrent_readers_of_different_datasets_see_their_own_chunks() {
        let cache = std::sync::Arc::new(ChunkCache::with_capacity(1 << 20, 64));
        let handles: Vec<_> = (0..8u64)
            .map(|t| {
                let cache = std::sync::Arc::clone(&cache);
                std::thread::spawn(move || {
                    for round in 0..500u64 {
                        let addr = (t + round) % 16;
                        let coord = vec![round % 4];
                        let chunks = cache
                            .chunks_for::<()>(addr, 1, || {
                                Ok((0..4).map(|c| make_chunk(vec![c], addr, 8)).collect())
                            })
                            .unwrap();
                        assert!(chunks.iter().all(|c| c.address == addr));
                        let want = vec![addr as u8; 8];
                        let got = match cache.get_decompressed_in(addr, &coord) {
                            Some(hit) => hit.to_vec(),
                            None => cache
                                .put_decompressed_in(addr, coord, want.clone())
                                .to_vec(),
                        };
                        assert_eq!(got, want);
                    }
                })
            })
            .collect();
        for h in handles {
            h.join().unwrap();
        }
    }

    #[test]
    fn duplicate_insert_is_noop() {
        let cache = ChunkCache::new();
        cache.put_decompressed(vec![0], vec![1, 2, 3]);
        cache.put_decompressed(vec![0], vec![1, 2, 3]); // duplicate
        assert_eq!(cache.cached_chunk_count(), 1);
        assert_eq!(cache.cached_bytes(), 3);
    }

    #[test]
    fn slot_index_consistent_after_many_evictions() {
        // Force repeated swap_remove evictions (small slot budget, many
        // inserts) and confirm the coord -> slot index stays correct: every
        // remaining coord must still resolve to its own data, not another
        // slot's (which would happen if swap_remove's index fixup were wrong).
        let cache = ChunkCache::with_capacity(1024 * 1024, 4); // max 4 slots

        for i in 0..50u64 {
            cache.put_decompressed(vec![i], vec![(i % 256) as u8; 8]);
            // Interleave reads of a couple of earlier coords to churn LRU
            // order (and thus which slot gets swap_remove'd) beyond simple
            // FIFO eviction.
            if i >= 2 {
                let _ = cache.get_decompressed(&[i - 2]);
            }
        }

        // Whatever remains in the cache (at most 4 slots) must return its
        // own correct data.
        for i in 0..50u64 {
            if let Some(data) = cache.get_decompressed(&[i]) {
                assert_eq!(
                    data,
                    vec![(i % 256) as u8; 8],
                    "coord {i} returned wrong data after eviction churn"
                );
            }
        }
        assert!(cache.cached_chunk_count() <= 4);
    }

    #[test]
    fn get_decompressed_aligned_shares_arc_on_hit() {
        let cache = ChunkCache::new();
        cache.put_decompressed(vec![0, 0], vec![9, 9, 9, 9]);
        let a = cache.get_decompressed_aligned(&[0, 0]).unwrap();
        let b = cache.get_decompressed_aligned(&[0, 0]).unwrap();
        // A cache hit clones the Arc (refcount bump), not the underlying
        // buffer — both handles point at the same allocation.
        assert!(Arc::ptr_eq(&a, &b));
        assert_eq!(a.as_slice(), &[9, 9, 9, 9]);
    }

    // --- CacheAlignedBuffer tests ---

    #[test]
    fn aligned_buffer_basic() {
        let buf = CacheAlignedBuffer::zeroed(256);
        assert_eq!(buf.len(), 256);
        assert!(buf.is_aligned());
        assert_eq!(&buf[..4], &[0, 0, 0, 0]);
    }

    #[test]
    fn aligned_buffer_from_slice() {
        let data = vec![1u8, 2, 3, 4, 5];
        let buf = CacheAlignedBuffer::from_slice(&data);
        assert_eq!(buf.len(), 5);
        assert!(buf.is_aligned());
        assert_eq!(buf.to_vec(), data);
    }

    #[test]
    fn aligned_buffer_from_vec() {
        let data = vec![42u8; 1024];
        let buf = CacheAlignedBuffer::from_vec(data.clone());
        assert!(buf.is_aligned());
        assert_eq!(buf.to_vec(), data);
    }

    #[test]
    fn aligned_buffer_empty() {
        let buf = CacheAlignedBuffer::zeroed(0);
        assert!(buf.is_empty());
        assert!(buf.is_aligned());
        assert_eq!(buf.to_vec(), Vec::<u8>::new());
    }

    #[test]
    fn aligned_buffer_clone_is_aligned() {
        let buf = CacheAlignedBuffer::from_slice(&[1, 2, 3, 4]);
        let cloned = buf.clone();
        assert!(cloned.is_aligned());
        assert_eq!(buf.to_vec(), cloned.to_vec());
    }

    #[test]
    fn aligned_buffer_deref_works() {
        let buf = CacheAlignedBuffer::from_slice(&[10, 20, 30]);
        assert_eq!(buf[0], 10);
        assert_eq!(buf[1], 20);
        assert_eq!(buf[2], 30);
    }

    #[test]
    fn aligned_buffer_various_sizes() {
        // Test alignment for various sizes including non-power-of-two
        for size in [1, 7, 63, 64, 65, 127, 128, 129, 255, 256, 1000, 4096] {
            let buf = CacheAlignedBuffer::zeroed(size);
            assert!(buf.is_aligned(), "not aligned for size {size}");
            assert_eq!(buf.len(), size);
        }
    }

    #[test]
    fn cached_data_is_aligned() {
        let cache = ChunkCache::new();
        cache.put_decompressed(vec![0, 0], vec![1, 2, 3, 4, 5, 6, 7, 8]);
        let aligned = cache.get_decompressed_aligned(&[0, 0]).unwrap();
        assert!(aligned.is_aligned());
        assert_eq!(aligned.to_vec(), vec![1, 2, 3, 4, 5, 6, 7, 8]);
    }

    #[test]
    fn align_to_cache_line_values() {
        assert_eq!(align_to_cache_line(0), 0);
        assert_eq!(align_to_cache_line(1), CACHE_LINE_SIZE);
        assert_eq!(align_to_cache_line(CACHE_LINE_SIZE), CACHE_LINE_SIZE);
        assert_eq!(
            align_to_cache_line(CACHE_LINE_SIZE + 1),
            CACHE_LINE_SIZE * 2
        );
    }
}