wsi-rs 0.5.2

use lru::LruCache;
use std::borrow::Borrow;
use std::hash::Hash;
use std::num::NonZeroUsize;
use std::sync::{Arc, Mutex};

use crate::core::types::{CpuTile, DatasetId};

// ── TileCache (axis-aware) ────────────────────────────────────────

/// Default shared decoded tile cache.
///
/// Standard Aperio SVS JPEG tiles are commonly 240x240 RGB, or about 170 KiB
/// per decoded tile. A 64 MiB budget keeps a few hundred such source tiles
/// resident, which is enough for normal viewport overlap during quick zooms
/// without forcing users to tune cache options before the viewer is usable.
pub(crate) const DEFAULT_TILE_CACHE_SIZE: u64 = 64 * 1024 * 1024;
const TILE_CACHE_BYTES_ENV: &str = "WSI_RS_TILE_CACHE_BYTES";
/// Default display-tile cache.
///
/// Display-tile reads on regular tiled slides cache the decoded source tiles
/// used for composition. Keep enough room for at least a dense viewport plus
/// adjacent zoom/pan overlap; 1 MiB only held a handful of SVS tiles and caused
/// immediate churn during zoom-out bursts.
pub(crate) const DEFAULT_DISPLAY_TILE_CACHE_SIZE: u64 = 32 * 1024 * 1024;
const DISPLAY_TILE_CACHE_BYTES_ENV: &str = "WSI_RS_DISPLAY_TILE_CACHE_BYTES";
// Count-bounded private LRUs still allocate per-entry bookkeeping. Include a
// conservative floor so tiny source tiles cannot turn a byte budget into an
// enormous hash-table capacity at open time.
const PRIVATE_CACHE_ENTRY_ACCOUNTING_FLOOR_BYTES: u64 = 256;
// Private per-format caches supplement the public shared tile cache. Bound
// their aggregate entry capacity to one quarter of the configured shared
// cache so opening a slide cannot multiply the caller's byte policy.
const PRIVATE_CACHE_BUDGET_DIVISOR: u64 = 4;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub struct CacheConfig {
    pub shared_tile_bytes: Option<u64>,
    pub display_tile_bytes: Option<u64>,
}

impl CacheConfig {
    pub const fn deterministic() -> Self {
        Self {
            shared_tile_bytes: None,
            display_tile_bytes: None,
        }
    }

    pub const fn with_shared_tile_bytes(mut self, bytes: u64) -> Self {
        self.shared_tile_bytes = Some(bytes);
        self
    }

    pub const fn with_display_tile_bytes(mut self, bytes: u64) -> Self {
        self.display_tile_bytes = Some(bytes);
        self
    }

    pub(crate) fn shared_tile_budget(self, source_hint: Option<u64>) -> u64 {
        self.shared_tile_bytes
            .or(source_hint)
            .unwrap_or(DEFAULT_TILE_CACHE_SIZE)
    }

    pub(crate) fn display_tile_budget(self) -> u64 {
        self.display_tile_bytes
            .unwrap_or(DEFAULT_DISPLAY_TILE_CACHE_SIZE)
    }

    pub(crate) fn private_cache_budget(self, cache_count: usize) -> PrivateCacheBudget {
        PrivateCacheBudget {
            remaining_bytes: self.private_cache_budget_bytes(),
            remaining_caches: cache_count,
        }
    }

    pub(crate) fn private_cache_budget_bytes(self) -> u64 {
        self.shared_tile_budget(None) / PRIVATE_CACHE_BUDGET_DIVISOR
    }
}

impl Default for CacheConfig {
    fn default() -> Self {
        Self::deterministic()
    }
}

/// One slide's aggregate budget for count-bounded format-private caches.
///
/// Allocations are made against estimated retained bytes (including a floor
/// for LRU bookkeeping). A cache receives zero entries when the remaining
/// budget cannot account for one entry, avoiding eager hash-table allocation.
pub(crate) struct PrivateCacheBudget {
    remaining_bytes: u64,
    remaining_caches: usize,
}

impl PrivateCacheBudget {
    pub(crate) fn allocate(&mut self, estimated_entry_bytes: u64) -> PrivateCacheCapacity {
        if self.remaining_caches == 0 {
            return PrivateCacheCapacity::default();
        }

        let cache_count = self.remaining_caches as u64;
        self.remaining_caches -= 1;
        let accounted_entry_bytes =
            estimated_entry_bytes.max(PRIVATE_CACHE_ENTRY_ACCOUNTING_FLOOR_BYTES);
        let fair_share = self.remaining_bytes / cache_count;
        let mut entries = fair_share / accounted_entry_bytes;
        if entries == 0 && self.remaining_bytes >= accounted_entry_bytes {
            entries = 1;
        }
        entries = entries.min(usize::MAX as u64);
        let accounted_bytes = entries
            .checked_mul(accounted_entry_bytes)
            .unwrap_or(self.remaining_bytes)
            .min(self.remaining_bytes);
        self.remaining_bytes -= accounted_bytes;

        PrivateCacheCapacity {
            entries: entries as usize,
            accounted_bytes,
        }
    }
}

#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub(crate) struct PrivateCacheCapacity {
    entries: usize,
    accounted_bytes: u64,
}

/// A count-bounded private LRU that can be completely disabled.
///
/// `lru::LruCache::new` preallocates for its entry capacity and requires a
/// non-zero value. Keeping the disabled state as `None` makes a zero-budget
/// cache allocation-free and causes inserts to be ignored.
#[derive(Debug)]
pub(crate) struct PrivateCache<K: Hash + Eq, V> {
    lru: Option<LruCache<K, V>>,
    capacity: PrivateCacheCapacity,
}

impl<K: Hash + Eq, V> PrivateCache<K, V> {
    pub(crate) fn new(capacity: PrivateCacheCapacity) -> Self {
        let lru = NonZeroUsize::new(capacity.entries).map(LruCache::new);
        Self { lru, capacity }
    }

    pub(crate) fn get<'a, Q>(&'a mut self, key: &Q) -> Option<&'a V>
    where
        K: Borrow<Q>,
        Q: Hash + Eq + ?Sized,
    {
        self.lru.as_mut()?.get(key)
    }

    pub(crate) fn put(&mut self, key: K, value: V) {
        if let Some(lru) = &mut self.lru {
            lru.put(key, value);
        }
    }

    pub(crate) fn capacity_entries(&self) -> usize {
        self.capacity.entries
    }

    #[cfg(test)]
    pub(crate) fn accounted_capacity_bytes(&self) -> u64 {
        self.capacity.accounted_bytes
    }

    #[cfg(test)]
    pub(crate) fn len(&self) -> usize {
        self.lru.as_ref().map_or(0, LruCache::len)
    }
}

#[derive(Hash, Eq, PartialEq, Clone, Debug)]
/// Note: scene/series are u32 here (not usize) to keep CacheKey compact and
/// Hash-friendly. TileRequest/RegionRequest use usize for ergonomic indexing.
/// Slide converts usize → u32 via `as u32` when constructing cache keys.
/// Overflow is not a practical concern (>4B scenes/series is impossible).
pub struct CacheKey {
    pub(crate) dataset_id: DatasetId,
    pub(crate) scene: u32,
    pub(crate) series: u32,
    pub(crate) level: u32,
    pub(crate) z: u32,
    pub(crate) c: u32,
    pub(crate) t: u32,
    pub(crate) tile_col: i64,
    pub(crate) tile_row: i64,
}

/// Thread-safe, byte-bounded decoded tile cache that can be shared by slides.
pub struct TileCache {
    inner: Mutex<TileCacheState>,
}

/// Snapshot of byte-sized decoded tile cache activity.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct TileCacheStats {
    /// Successful lookups.
    pub hits: u64,
    /// Unsuccessful lookups.
    pub misses: u64,
    /// Entries admitted to the cache.
    pub puts: u64,
    /// Entries removed to remain within the byte capacity.
    pub evictions: u64,
    /// Entries rejected because one value exceeded the whole capacity.
    pub rejected_oversize: u64,
    /// Configured byte capacity.
    pub capacity_bytes: u64,
    /// Bytes currently retained by cached entries.
    pub current_bytes: u64,
    /// Number of entries currently retained.
    pub entries: usize,
}

impl std::fmt::Debug for TileCache {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let state = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        f.debug_struct("TileCache")
            .field("capacity_bytes", &state.capacity_bytes)
            .field("current_bytes", &state.current_bytes)
            .field("entries", &state.lru.len())
            .field("hits", &state.hits)
            .field("misses", &state.misses)
            .finish()
    }
}

struct TileCacheState {
    lru: LruCache<CacheKey, CachedTile>,
    capacity_bytes: u64,
    current_bytes: u64,
    hits: u64,
    misses: u64,
    puts: u64,
    evictions: u64,
    rejected_oversize: u64,
}

struct CachedTile {
    data: Arc<CpuTile>,
    byte_size: u64,
}

impl TileCache {
    /// Create a thread-safe decoded tile cache with a byte capacity.
    pub fn new(capacity_bytes: u64) -> Self {
        Self {
            inner: Mutex::new(TileCacheState {
                // The cache is byte-budgeted only. The backing LRU stays unbounded
                // and eviction is driven by `capacity_bytes`.
                lru: LruCache::unbounded(),
                capacity_bytes,
                current_bytes: 0,
                hits: 0,
                misses: 0,
                puts: 0,
                evictions: 0,
                rejected_oversize: 0,
            }),
        }
    }

    pub(crate) fn put(&self, key: CacheKey, data: Arc<CpuTile>) {
        let byte_size = data.data.byte_size() as u64;
        let mut state = self.inner.lock().unwrap_or_else(|e| e.into_inner());

        if byte_size > state.capacity_bytes {
            state.rejected_oversize += 1;
            return;
        }

        // Remove existing entry if present
        if let Some((_, existing)) = state.lru.pop_entry(&key) {
            state.current_bytes -= existing.byte_size;
        }

        // Evict LRU entries until there's room
        while state.current_bytes + byte_size > state.capacity_bytes {
            if let Some((_, evicted)) = state.lru.pop_lru() {
                state.current_bytes -= evicted.byte_size;
                state.evictions += 1;
            } else {
                break;
            }
        }

        state.lru.put(key, CachedTile { data, byte_size });
        state.current_bytes += byte_size;
        state.puts += 1;
    }

    pub(crate) fn get(&self, key: &CacheKey) -> Option<Arc<CpuTile>> {
        let mut state = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        let cached = state.lru.get(key).map(|entry| entry.data.clone());
        if cached.is_some() {
            state.hits += 1;
        } else {
            state.misses += 1;
        }
        cached
    }

    /// Return an atomic snapshot of cache capacity and activity counters.
    pub fn stats(&self) -> TileCacheStats {
        let state = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        TileCacheStats {
            hits: state.hits,
            misses: state.misses,
            puts: state.puts,
            evictions: state.evictions,
            rejected_oversize: state.rejected_oversize,
            capacity_bytes: state.capacity_bytes,
            current_bytes: state.current_bytes,
            entries: state.lru.len(),
        }
    }

    pub(crate) fn display_default() -> Self {
        Self::new(capacity_from_env(
            DISPLAY_TILE_CACHE_BYTES_ENV,
            DEFAULT_DISPLAY_TILE_CACHE_SIZE,
        ))
    }

    pub(crate) fn display_with_config(config: CacheConfig) -> Self {
        Self::new(config.display_tile_budget())
    }

    pub(crate) fn shared_default_with_hint(default_bytes: u64) -> Self {
        Self::new(capacity_from_env(TILE_CACHE_BYTES_ENV, default_bytes))
    }

    pub(crate) fn shared_with_config(config: CacheConfig, source_hint: Option<u64>) -> Self {
        Self::new(config.shared_tile_budget(source_hint))
    }
}

impl Default for TileCache {
    fn default() -> Self {
        Self::shared_default_with_hint(DEFAULT_TILE_CACHE_SIZE)
    }
}

fn capacity_from_env(env_name: &str, default_bytes: u64) -> u64 {
    std::env::var(env_name)
        .ok()
        .and_then(|value| value.parse::<u64>().ok())
        .filter(|bytes| *bytes > 0)
        .unwrap_or(default_bytes)
}

#[cfg(test)]
mod tile_cache_tests {
    use super::*;
    use crate::core::types::*;

    const SVS_RGB_240_TILE_BYTES: usize = 240 * 240 * 3;
    const COMMON_ZOOM_VIEWPORT_TILE_COUNT: i64 = 96;

    fn make_sample_buffer(size: usize) -> CpuTile {
        CpuTile {
            width: 256,
            height: 256,
            channels: 3,
            color_space: ColorSpace::Rgb,
            layout: CpuTileLayout::Interleaved,
            data: CpuTileData::u8(vec![0u8; size]),
        }
    }

    fn make_key(dataset_id: u128, level: u32, col: i64, row: i64) -> CacheKey {
        CacheKey {
            dataset_id: DatasetId::new(dataset_id),
            scene: 0,
            series: 0,
            level,
            z: 0,
            c: 0,
            t: 0,
            tile_col: col,
            tile_row: row,
        }
    }

    #[test]
    fn private_cache_budget_bounds_aggregate_capacity_and_disables_excess_caches() {
        let config = CacheConfig::deterministic().with_shared_tile_bytes(8 * 1024);
        let mut budget = config.private_cache_budget(8);
        let mut caches = (0..8)
            .map(|_| PrivateCache::<u32, u32>::new(budget.allocate(1024)))
            .collect::<Vec<_>>();

        assert!(
            caches
                .iter()
                .map(PrivateCache::accounted_capacity_bytes)
                .sum::<u64>()
                <= config.private_cache_budget_bytes()
        );
        assert_eq!(
            caches
                .iter()
                .map(PrivateCache::capacity_entries)
                .sum::<usize>(),
            2
        );
        assert_eq!(
            caches
                .iter()
                .filter(|cache| cache.capacity_entries() == 0)
                .count(),
            6
        );

        let disabled = caches
            .iter_mut()
            .find(|cache| cache.capacity_entries() == 0)
            .expect("small aggregate budget disables at least one cache");
        disabled.put(1, 2);
        assert_eq!(disabled.len(), 0, "disabled caches retain no entries");
    }

    #[test]
    fn put_and_get() {
        let cache = TileCache::new(1024 * 1024);
        let buf = Arc::new(make_sample_buffer(100));
        let key = make_key(1, 0, 0, 0);
        cache.put(key.clone(), buf.clone());
        let result = cache.get(&key).unwrap();
        assert_eq!(result.width, 256);
    }

    #[test]
    fn miss_returns_none() {
        let cache = TileCache::new(1024);
        let key = make_key(1, 0, 0, 0);
        assert!(cache.get(&key).is_none());
    }

    #[test]
    fn eviction_by_byte_size() {
        let cache = TileCache::new(250);
        cache.put(make_key(1, 0, 0, 0), Arc::new(make_sample_buffer(100)));
        cache.put(make_key(1, 0, 1, 0), Arc::new(make_sample_buffer(100)));
        // Both fit: 200 bytes
        assert!(cache.get(&make_key(1, 0, 0, 0)).is_some());
        assert!(cache.get(&make_key(1, 0, 1, 0)).is_some());

        // Third pushes over 250
        cache.put(make_key(1, 0, 2, 0), Arc::new(make_sample_buffer(100)));
        assert!(cache.get(&make_key(1, 0, 0, 0)).is_none()); // evicted
        assert!(cache.get(&make_key(1, 0, 1, 0)).is_some());
        assert!(cache.get(&make_key(1, 0, 2, 0)).is_some());
    }

    #[test]
    fn different_datasets_are_independent() {
        let cache = TileCache::new(1024);
        cache.put(make_key(1, 0, 0, 0), Arc::new(make_sample_buffer(10)));
        cache.put(make_key(2, 0, 0, 0), Arc::new(make_sample_buffer(10)));
        assert!(cache.get(&make_key(1, 0, 0, 0)).is_some());
        assert!(cache.get(&make_key(2, 0, 0, 0)).is_some());
    }

    #[test]
    fn axis_aware_keys() {
        let cache = TileCache::new(1024);
        let mut key_z0 = make_key(1, 0, 0, 0);
        key_z0.z = 0;
        let mut key_z1 = make_key(1, 0, 0, 0);
        key_z1.z = 1;
        cache.put(key_z0.clone(), Arc::new(make_sample_buffer(10)));
        cache.put(key_z1.clone(), Arc::new(make_sample_buffer(10)));
        assert!(cache.get(&key_z0).is_some());
        assert!(cache.get(&key_z1).is_some());
    }

    #[test]
    fn oversize_entry_rejected() {
        let cache = TileCache::new(50);
        cache.put(make_key(1, 0, 0, 0), Arc::new(make_sample_buffer(100)));
        assert!(cache.get(&make_key(1, 0, 0, 0)).is_none());
    }

    #[test]
    fn shared_across_threads() {
        let cache = Arc::new(TileCache::new(4096));
        let cache_clone = cache.clone();
        let handle = std::thread::spawn(move || {
            cache_clone.put(make_key(1, 0, 5, 5), Arc::new(make_sample_buffer(10)));
        });
        handle.join().unwrap();
        assert!(cache.get(&make_key(1, 0, 5, 5)).is_some());
    }

    #[test]
    fn display_default_holds_common_svs_zoom_viewport_working_set() {
        let cache = TileCache::new(DEFAULT_DISPLAY_TILE_CACHE_SIZE);
        for col in 0..COMMON_ZOOM_VIEWPORT_TILE_COUNT {
            cache.put(
                make_key(1, 0, col, 0),
                Arc::new(make_sample_buffer(SVS_RGB_240_TILE_BYTES)),
            );
        }

        let stats = cache.stats();
        assert_eq!(stats.entries, COMMON_ZOOM_VIEWPORT_TILE_COUNT as usize);
        assert_eq!(stats.evictions, 0);
        assert_eq!(stats.rejected_oversize, 0);
    }

    #[test]
    fn stats_count_hits_misses_puts_evictions_and_oversize_rejections() {
        let cache = TileCache::new(150);
        let missing = make_key(1, 0, 9, 9);
        assert!(cache.get(&missing).is_none());

        cache.put(make_key(1, 0, 0, 0), Arc::new(make_sample_buffer(100)));
        assert!(cache.get(&make_key(1, 0, 0, 0)).is_some());

        cache.put(make_key(1, 0, 1, 0), Arc::new(make_sample_buffer(100)));
        cache.put(make_key(1, 0, 2, 0), Arc::new(make_sample_buffer(200)));

        let stats = cache.stats();
        assert_eq!(stats.hits, 1);
        assert_eq!(stats.misses, 1);
        assert_eq!(stats.puts, 2);
        assert_eq!(stats.evictions, 1);
        assert_eq!(stats.rejected_oversize, 1);
        assert_eq!(stats.capacity_bytes, 150);
        assert_eq!(stats.current_bytes, 100);
        assert_eq!(stats.entries, 1);
    }
}