jellytau/src-tauri/src/download/cache.rs

//! Smart caching engine for predictive downloads

#[cfg(test)]
use crate::utils::lock::MutexSafe;
use log::{debug, info};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use serde::{Deserialize, Serialize};

use crate::storage::db_service::{DatabaseService, Query, QueryParam};

/// Smart caching configuration
#[derive(specta::Type, Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CacheConfig {
    /// Enable queue pre-caching
    pub queue_precache_enabled: bool,
    /// Number of tracks to pre-cache from queue
    pub queue_precache_count: usize,
    /// Enable album affinity detection
    pub album_affinity_enabled: bool,
    /// Threshold for album affinity (tracks played before caching)
    pub album_affinity_threshold: usize,
    /// Storage limit in bytes (0 = unlimited)
    pub storage_limit: u64,
    /// Only cache on WiFi
    pub wifi_only: bool,
}

impl Default for CacheConfig {
    fn default() -> Self {
        Self {
            queue_precache_enabled: true,
            queue_precache_count: 3, // Preload next 3 tracks by default
            album_affinity_enabled: true,
            album_affinity_threshold: 3,
            storage_limit: 10 * 1024 * 1024 * 1024, // 10GB
            wifi_only: false, // Allow preloading on any connection by default
        }
    }
}

/// Smart caching engine
#[derive(Clone)]
pub struct SmartCache {
    config: Arc<Mutex<CacheConfig>>,
    /// Track recently played items per album
    album_play_history: Arc<Mutex<HashMap<String, Vec<String>>>>,
}

impl SmartCache {
    pub fn new(config: CacheConfig) -> Self {
        Self {
            config: Arc::new(Mutex::new(config)),
            album_play_history: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// Update configuration
    pub fn update_config(&self, config: CacheConfig) {
        if let Ok(mut cfg) = self.config.lock() {
            *cfg = config;
        }
    }

    /// Check if should pre-cache queue items
    pub fn should_precache_queue(&self) -> bool {
        self.config
            .lock()
            .map(|cfg| cfg.queue_precache_enabled && !cfg.wifi_only)
            .unwrap_or(false)
    }

    /// Get number of queue items to pre-cache
    pub fn queue_precache_count(&self) -> usize {
        self.config
            .lock()
            .map(|cfg| cfg.queue_precache_count)
            .unwrap_or(5)
    }

    /// Track that an item was played
    pub fn track_play(&self, item_id: &str, album_id: Option<&str>) {
        if let Some(album) = album_id {
            if let Ok(mut history) = self.album_play_history.lock() {
                let plays = history.entry(album.to_string()).or_insert_with(Vec::new);
                if !plays.contains(&item_id.to_string()) {
                    plays.push(item_id.to_string());
                }
            }
        }
    }

    /// Check if album affinity threshold reached for caching
    pub fn should_cache_album(&self, album_id: &str) -> Option<bool> {
        let config = self.config.lock().ok()?;
        if !config.album_affinity_enabled {
            return Some(false);
        }

        let history = self.album_play_history.lock().ok()?;
        let play_count = history.get(album_id).map(|v| v.len()).unwrap_or(0);

        Some(play_count >= config.album_affinity_threshold)
    }

    /// Get configuration
    pub fn get_config(&self) -> Option<CacheConfig> {
        self.config.lock().ok().map(|cfg| cfg.clone())
    }

    /// Get all tracked albums with their play counts
    /// Returns Vec<(album_id, unique_tracks_played)>
    pub fn get_album_play_history(&self) -> Vec<(String, usize)> {
        self.album_play_history
            .lock()
            .ok()
            .map(|history| {
                history
                    .iter()
                    .map(|(album_id, tracks)| (album_id.clone(), tracks.len()))
                    .collect()
            })
            .unwrap_or_default()
    }

    // ============= Async versions for DatabaseService =============

    /// Get total download size for a user (async version)
    pub async fn get_total_download_size_async<S: DatabaseService>(
        &self,
        db_service: &Arc<S>,
        user_id: &str,
    ) -> Result<u64, String> {
        let query = Query::with_params(
            "SELECT COALESCE(SUM(file_size), 0) FROM downloads
             WHERE user_id = ? AND status = 'completed'",
            vec![QueryParam::String(user_id.to_string())],
        );

        let size: i64 = db_service
            .query_one(query, |row| row.get(0))
            .await
            .map_err(|e| e.to_string())?;

        Ok(size as u64)
    }

    /// Check if storage limit allows download (async version)
    pub async fn can_download_async<S: DatabaseService>(
        &self,
        db_service: &Arc<S>,
        user_id: &str,
        new_size: u64,
    ) -> bool {
        // Clone config to avoid holding lock across await
        let storage_limit = {
            match self.config.lock() {
                Ok(cfg) => cfg.storage_limit,
                Err(_) => return true,
            }
        };

        if storage_limit == 0 {
            return true; // Unlimited
        }

        let current_size = self
            .get_total_download_size_async(db_service, user_id)
            .await
            .unwrap_or(0);

        current_size + new_size <= storage_limit
    }

    /// Evict least recently used items to make space (async version)
    pub async fn evict_lru_async<S: DatabaseService>(
        &self,
        db_service: &Arc<S>,
        user_id: &str,
        space_needed: u64,
    ) -> Result<u64, String> {
        let current_size = self
            .get_total_download_size_async(db_service, user_id)
            .await?;

        // Get limit without holding lock across await
        let limit = {
            let config = self.config.lock().map_err(|e| e.to_string())?;
            config.storage_limit
        };

        if limit == 0 || current_size + space_needed <= limit {
            return Ok(0); // No eviction needed
        }

        let to_free = (current_size + space_needed) - limit;
        let mut freed: u64 = 0;

        // Get downloads ordered by last access (oldest first)
        let query = Query::with_params(
            "SELECT id, file_size, file_path FROM downloads
             WHERE user_id = ? AND status = 'completed'
             ORDER BY completed_at ASC",
            vec![QueryParam::String(user_id.to_string())],
        );

        let downloads: Vec<(i64, i64, String)> = db_service
            .query_many(query, |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)))
            .await
            .map_err(|e| e.to_string())?;

        for (id, size, file_path) in downloads {
            if freed >= to_free {
                break;
            }

            // Delete file
            let _ = std::fs::remove_file(&file_path);
            debug!("[SmartCache] Evicted: {} ({} bytes)", file_path, size);

            // Delete from database
            let delete_query = Query::with_params(
                "DELETE FROM downloads WHERE id = ?",
                vec![QueryParam::Int64(id)],
            );
            db_service.execute(delete_query).await.map_err(|e| e.to_string())?;

            freed += size as u64;
        }

        info!("[SmartCache] Freed {} bytes ({} needed)", freed, to_free);
        Ok(freed)
    }
}

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

    #[test]
    fn test_default_config() {
        let config = CacheConfig::default();
        assert_eq!(config.queue_precache_count, 3);
        assert_eq!(config.album_affinity_threshold, 3);
        assert!(!config.wifi_only); // wifi_only is false by default for easier preloading
    }

    #[test]
    fn test_album_affinity_tracking() {
        let cache = SmartCache::new(CacheConfig::default());

        // Track plays from same album
        cache.track_play("track1", Some("album1"));
        cache.track_play("track2", Some("album1"));

        // Below threshold
        assert!(!cache.should_cache_album("album1").unwrap_or(false));

        cache.track_play("track3", Some("album1"));

        // At threshold - should cache
        assert!(cache.should_cache_album("album1").unwrap_or(false));
    }

    #[test]
    fn test_queue_precache_config() {
        let mut config = CacheConfig::default();
        config.queue_precache_enabled = false;

        let cache = SmartCache::new(config);
        assert!(!cache.should_precache_queue());

        let mut new_config = CacheConfig::default();
        new_config.wifi_only = false;
        cache.update_config(new_config);

        assert!(cache.should_precache_queue());
    }

    #[tokio::test]
    async fn test_storage_limit_check() {
        use crate::storage::db_service::RusqliteService;
        use rusqlite::Connection;
        use std::sync::{Arc, Mutex};

        let conn = Connection::open_in_memory().unwrap();
        conn.execute(
            "CREATE TABLE downloads (
                id INTEGER PRIMARY KEY,
                user_id TEXT,
                status TEXT,
                file_size INTEGER
            )",
            [],
        )
        .unwrap();

        let conn_arc = Arc::new(Mutex::new(conn));
        let db_service = Arc::new(RusqliteService::new(conn_arc.clone()));

        let config = CacheConfig {
            storage_limit: 1000,
            ..Default::default()
        };
        let cache = SmartCache::new(config);

        // Empty - can download
        assert!(cache.can_download_async(&db_service, "user1", 500).await);

        // Add some downloads
        {
            let conn_guard = conn_arc.lock_safe();
            conn_guard.execute(
                "INSERT INTO downloads (user_id, status, file_size) VALUES ('user1', 'completed', 600)",
                [],
            )
            .unwrap();
        }

        // Total would be 1100 > 1000
        assert!(!cache.can_download_async(&db_service, "user1", 500).await);

        // Smaller size fits
        assert!(cache.can_download_async(&db_service, "user1", 300).await);
    }
}