ai-lib-contact 0.9.6

//! Cache manager.

use super::backend::CacheBackend;
use super::key::CacheKey;
use ai_lib_core::Result;
use serde::{de::DeserializeOwned, Serialize};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;

#[derive(Debug, Clone)]
pub struct CacheConfig {
    pub default_ttl: Duration,
    pub enabled: bool,
    pub max_entry_size: usize,
    pub key_prefix: Option<String>,
}

impl Default for CacheConfig {
    fn default() -> Self {
        Self {
            default_ttl: Duration::from_secs(3600),
            enabled: true,
            max_entry_size: 10 * 1024 * 1024,
            key_prefix: None,
        }
    }
}

impl CacheConfig {
    pub fn new() -> Self {
        Self::default()
    }
    pub fn with_ttl(mut self, ttl: Duration) -> Self {
        self.default_ttl = ttl;
        self
    }
    pub fn with_enabled(mut self, enabled: bool) -> Self {
        self.enabled = enabled;
        self
    }
    pub fn with_key_prefix(mut self, prefix: impl Into<String>) -> Self {
        self.key_prefix = Some(prefix.into());
        self
    }
}

#[derive(Debug, Clone, Default)]
pub struct CacheStats {
    pub hits: u64,
    pub misses: u64,
    pub sets: u64,
    pub deletes: u64,
    pub errors: u64,
}

impl CacheStats {
    pub fn hit_ratio(&self) -> f64 {
        let total = self.hits + self.misses;
        if total == 0 {
            0.0
        } else {
            self.hits as f64 / total as f64
        }
    }
}

struct AtomicStats {
    hits: AtomicU64,
    misses: AtomicU64,
    sets: AtomicU64,
    deletes: AtomicU64,
    errors: AtomicU64,
}
impl AtomicStats {
    fn new() -> Self {
        Self {
            hits: AtomicU64::new(0),
            misses: AtomicU64::new(0),
            sets: AtomicU64::new(0),
            deletes: AtomicU64::new(0),
            errors: AtomicU64::new(0),
        }
    }
    fn to_stats(&self) -> CacheStats {
        CacheStats {
            hits: self.hits.load(Ordering::Relaxed),
            misses: self.misses.load(Ordering::Relaxed),
            sets: self.sets.load(Ordering::Relaxed),
            deletes: self.deletes.load(Ordering::Relaxed),
            errors: self.errors.load(Ordering::Relaxed),
        }
    }
}

pub struct CacheManager {
    config: CacheConfig,
    backend: Box<dyn CacheBackend>,
    stats: Arc<AtomicStats>,
}

impl CacheManager {
    pub fn new(config: CacheConfig, backend: Box<dyn CacheBackend>) -> Self {
        Self {
            config,
            backend,
            stats: Arc::new(AtomicStats::new()),
        }
    }

    pub async fn get<T: DeserializeOwned>(&self, key: &CacheKey) -> Result<Option<T>> {
        if !self.config.enabled {
            return Ok(None);
        }
        let prefixed = self.prefix_key(key);
        match self.backend.get(&prefixed).await {
            Ok(Some(data)) => {
                self.stats.hits.fetch_add(1, Ordering::Relaxed);
                match serde_json::from_slice(&data) {
                    Ok(val) => Ok(Some(val)),
                    Err(_) => {
                        self.stats.errors.fetch_add(1, Ordering::Relaxed);
                        Ok(None)
                    }
                }
            }
            Ok(None) => {
                self.stats.misses.fetch_add(1, Ordering::Relaxed);
                Ok(None)
            }
            Err(e) => {
                self.stats.errors.fetch_add(1, Ordering::Relaxed);
                Err(e)
            }
        }
    }

    pub async fn set<T: Serialize>(&self, key: &CacheKey, value: &T) -> Result<()> {
        self.set_with_ttl(key, value, self.config.default_ttl).await
    }

    pub async fn set_with_ttl<T: Serialize>(
        &self,
        key: &CacheKey,
        value: &T,
        ttl: Duration,
    ) -> Result<()> {
        if !self.config.enabled {
            return Ok(());
        }
        let data = serde_json::to_vec(value)?;
        if data.len() > self.config.max_entry_size {
            return Ok(());
        }
        let prefixed = self.prefix_key(key);
        match self.backend.set(&prefixed, &data, ttl).await {
            Ok(()) => {
                self.stats.sets.fetch_add(1, Ordering::Relaxed);
                Ok(())
            }
            Err(e) => {
                self.stats.errors.fetch_add(1, Ordering::Relaxed);
                Err(e)
            }
        }
    }

    pub async fn delete(&self, key: &CacheKey) -> Result<bool> {
        if !self.config.enabled {
            return Ok(false);
        }
        let prefixed = self.prefix_key(key);
        match self.backend.delete(&prefixed).await {
            Ok(d) => {
                if d {
                    self.stats.deletes.fetch_add(1, Ordering::Relaxed);
                }
                Ok(d)
            }
            Err(e) => {
                self.stats.errors.fetch_add(1, Ordering::Relaxed);
                Err(e)
            }
        }
    }

    pub fn stats(&self) -> CacheStats {
        self.stats.to_stats()
    }
    pub fn backend_name(&self) -> &'static str {
        self.backend.name()
    }

    fn prefix_key(&self, key: &CacheKey) -> CacheKey {
        if let Some(ref p) = self.config.key_prefix {
            CacheKey::new(format!("{}:{}", p, key.hash))
        } else {
            key.clone()
        }
    }
}