baichun-framework-cache 0.1.0

mod config;
mod error;

use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::{OnceCell, RwLock};

use crate::{LocalClient, RedisClient};
pub use config::{CacheMode, MultilevelConfig};
pub use error::{MultilevelError, Result};

pub static MULTILEVEL_CACHE: OnceCell<Arc<MultilevelClient>> = OnceCell::const_new();

/// Circuit breaker state
#[derive(Debug, Clone)]
struct CircuitBreaker {
    /// Whether the circuit breaker is open (Redis disabled)
    is_open: bool,
    /// Time when the circuit breaker was opened
    opened_at: Option<Instant>,
    /// Count of consecutive failures
    failure_count: u32,
    /// Time of the first failure in current window
    first_failure: Option<Instant>,
}

impl Default for CircuitBreaker {
    fn default() -> Self {
        Self {
            is_open: false,
            opened_at: None,
            failure_count: 0,
            first_failure: None,
        }
    }
}

impl CircuitBreaker {
    fn reset(&mut self) {
        self.is_open = false;
        self.failure_count = 0;
        self.first_failure = None;
        self.opened_at = None;
    }
}

/// Multilevel cache client combining local and Redis caches
#[derive(Clone)]
pub struct MultilevelClient {
    local: Arc<LocalClient>,
    redis: Arc<RedisClient>,
    config: MultilevelConfig,
    circuit_breaker: Arc<RwLock<CircuitBreaker>>,
}

impl MultilevelClient {
    /// Create a new multilevel cache client
    pub fn new(local: Arc<LocalClient>, redis: Arc<RedisClient>, config: MultilevelConfig) -> Self {
        Self {
            local,
            redis,
            config,
            circuit_breaker: Arc::new(RwLock::new(CircuitBreaker::default())),
        }
    }

    /// Check if Redis is available based on circuit breaker state
    async fn is_redis_available(&self) -> bool {
        if !self.config.enable_circuit_breaker {
            return true;
        }

        let mut breaker = self.circuit_breaker.write().await;

        // If circuit breaker is open, check if we should attempt reset
        if breaker.is_open {
            if let Some(opened_at) = breaker.opened_at {
                if opened_at.elapsed() >= Duration::from_secs(self.config.reset_timeout_secs) {
                    log::info!("Attempting to reset circuit breaker after timeout");
                    breaker.reset();
                    // 在重置后立即返回 true，允许尝试 Redis 操作
                    return true;
                }
            }
            // 如果熔断器打开且未到重置时间，返回 false
            return false;
        }
        // 如果熔断器关闭，返回 true
        true
    }

    /// Record a Redis failure and possibly trigger circuit breaker
    async fn record_redis_failure(&self) {
        if !self.config.enable_circuit_breaker {
            return;
        }

        let mut breaker = self.circuit_breaker.write().await;
        let now = Instant::now();

        // Reset failure count if we're outside the failure window
        if let Some(first_failure) = breaker.first_failure {
            if first_failure.elapsed() >= Duration::from_secs(self.config.failure_window_secs) {
                breaker.reset();
                // 重置后立即返回，不增加失败计数
                return;
            }
        }

        // Record failure
        breaker.failure_count += 1;
        if breaker.first_failure.is_none() {
            breaker.first_failure = Some(now);
        }

        // Check if we should open the circuit breaker
        if breaker.failure_count >= self.config.failure_threshold {
            breaker.is_open = true;
            breaker.opened_at = Some(now);
            log::warn!(
                "Circuit breaker opened due to Redis failures: count={}",
                breaker.failure_count
            );
        }
    }

    /// Record a successful Redis operation
    async fn record_redis_success(&self) {
        if !self.config.enable_circuit_breaker {
            return;
        }

        let mut breaker = self.circuit_breaker.write().await;
        // 成功操作时总是重置熔断器状态
        breaker.reset();
        log::info!("Circuit breaker reset after successful Redis operation");
    }

    /// Set a key-value pair with automatic fallback to local cache if Redis fails
    pub async fn set<T: ToString + Clone>(&self, key: &str, value: T) -> Result<()> {
        match self.config.mode {
            CacheMode::Local => self
                .local
                .set(key, value)
                .await
                .map_err(MultilevelError::Local),
            CacheMode::Redis => {
                if self.is_redis_available().await {
                    match self.redis.set(key, value).await {
                        Ok(_) => {
                            self.record_redis_success().await;
                            Ok(())
                        }
                        Err(e) => {
                            self.record_redis_failure().await;
                            Err(MultilevelError::Redis(e))
                        }
                    }
                } else {
                    Err(MultilevelError::RedisUnavailable)
                }
            }
            CacheMode::Multi => {
                let mut redis_success = true;
                // Try Redis first if available
                if self.is_redis_available().await {
                    match self.redis.set(key, value.clone()).await {
                        Ok(_) => {
                            self.record_redis_success().await;
                        }
                        Err(e) => {
                            redis_success = false;
                            self.record_redis_failure().await;
                            log::warn!("Redis set failed, falling back to local cache: {}", e);
                        }
                    }
                } else {
                    redis_success = false;
                }

                // Always write to local cache in Multi mode if write_through is enabled or Redis failed
                if self.config.write_through || !redis_success {
                    self.local
                        .set(key, value)
                        .await
                        .map_err(MultilevelError::Local)?;
                }
                Ok(())
            }
        }
    }

    /// Get value by key
    pub async fn get(&self, key: &str) -> Result<Option<String>> {
        match self.config.mode {
            CacheMode::Local => self.local.get(key).await.map_err(MultilevelError::Local),
            CacheMode::Redis => {
                if self.is_redis_available().await {
                    match self.redis.get(key).await {
                        Ok(value) => {
                            self.record_redis_success().await;
                            Ok(value)
                        }
                        Err(e) => {
                            self.record_redis_failure().await;
                            Err(MultilevelError::Redis(e))
                        }
                    }
                } else {
                    Err(MultilevelError::RedisUnavailable)
                }
            }
            CacheMode::Multi => {
                // Try local cache first
                if let Some(value) = self.local.get(key).await.map_err(MultilevelError::Local)? {
                    return Ok(Some(value));
                }

                // If not in local cache and Redis is available, try Redis
                if self.is_redis_available().await {
                    match self.redis.get(key).await {
                        Ok(Some(value)) => {
                            self.record_redis_success().await;
                            // Update local cache if read-through is enabled
                            if self.config.read_through {
                                if let Err(e) = self.local.set(key, &value).await {
                                    log::warn!("Failed to read-through to local cache: {}", e);
                                }
                            }
                            Ok(Some(value))
                        }
                        Ok(None) => {
                            self.record_redis_success().await;
                            Ok(None)
                        }
                        Err(e) => {
                            self.record_redis_failure().await;
                            log::warn!("Redis get failed: {}", e);
                            // Return None since we already checked local cache
                            Ok(None)
                        }
                    }
                } else {
                    // Redis is unavailable, we already checked local cache
                    Ok(None)
                }
            }
        }
    }

    /// Delete a key
    pub async fn del(&self, key: &str) -> Result<bool> {
        match self.config.mode {
            CacheMode::Local => self.local.del(key).await.map_err(MultilevelError::Local),
            CacheMode::Redis => {
                if self.is_redis_available().await {
                    match self.redis.del(key).await {
                        Ok(deleted) => {
                            self.record_redis_success().await;
                            Ok(deleted)
                        }
                        Err(e) => {
                            self.record_redis_failure().await;
                            Err(MultilevelError::Redis(e))
                        }
                    }
                } else {
                    Err(MultilevelError::RedisUnavailable)
                }
            }
            CacheMode::Multi => {
                let mut deleted = false;

                // Try Redis first if available
                if self.is_redis_available().await {
                    match self.redis.del(key).await {
                        Ok(redis_deleted) => {
                            self.record_redis_success().await;
                            deleted = redis_deleted;
                        }
                        Err(e) => {
                            self.record_redis_failure().await;
                            log::warn!("Redis delete failed: {}", e);
                        }
                    }
                }

                // Always try to delete from local cache in Multi mode
                match self.local.del(key).await {
                    Ok(local_deleted) => deleted = deleted || local_deleted,
                    Err(e) => log::warn!("Local cache delete failed: {}", e),
                }

                Ok(deleted)
            }
        }
    }

    /// Clear all cache entries
    pub async fn clear(&self) -> Result<()> {
        match self.config.mode {
            CacheMode::Local => {
                self.local.clear().await;
                // 等待一小段时间确保清除完成
                tokio::time::sleep(Duration::from_millis(100)).await;
                Ok(())
            }
            CacheMode::Redis => {
                if self.is_redis_available().await {
                    // Note: Redis clear implementation needs to be added
                    Ok(())
                } else {
                    Err(MultilevelError::RedisUnavailable)
                }
            }
            CacheMode::Multi => {
                if self.is_redis_available().await {
                    // Try to clear Redis, but continue even if it fails
                    if let Err(e) = self.redis.clear().await {
                        self.record_redis_failure().await;
                        log::warn!("Failed to clear Redis cache: {}", e);
                    } else {
                        self.record_redis_success().await;
                    }
                }

                // Always clear local cache
                self.local.clear().await;
                Ok(())
            }
        }
    }

    pub async fn keys(&self, pattern: &str) -> Result<Vec<String>> {
        match self.config.mode {
            CacheMode::Local => self.local.keys(pattern).await.map_err(MultilevelError::Local),
            CacheMode::Redis => self.redis.keys(pattern).await.map_err(MultilevelError::Redis),
            CacheMode::Multi => self.redis.keys(pattern).await.map_err(MultilevelError::Redis),
        }
    }

    /// Get cache statistics
    pub fn stats(&self) -> MultilevelCacheStats {
        let local_stats = self.local.stats();
        MultilevelCacheStats {
            mode: self.config.mode,
            local_hits: local_stats.hits,
            local_misses: local_stats.misses,
            local_size: local_stats.size,
            circuit_breaker_enabled: self.config.enable_circuit_breaker,
        }
    }

    /// Get current circuit breaker status
    pub async fn circuit_breaker_status(&self) -> CircuitBreakerStatus {
        let breaker = self.circuit_breaker.read().await;
        CircuitBreakerStatus {
            enabled: self.config.enable_circuit_breaker,
            is_open: breaker.is_open,
            failure_count: breaker.failure_count,
            opened_at: breaker.opened_at,
        }
    }
}

/// Cache statistics
#[derive(Debug, Clone)]
pub struct MultilevelCacheStats {
    /// Current cache mode
    pub mode: CacheMode,
    /// Number of local cache hits
    pub local_hits: u64,
    /// Number of local cache misses
    pub local_misses: u64,
    /// Current local cache size
    pub local_size: u64,
    /// Whether circuit breaker is enabled
    pub circuit_breaker_enabled: bool,
}

/// Circuit breaker status
#[derive(Debug, Clone)]
pub struct CircuitBreakerStatus {
    /// Whether circuit breaker is enabled
    pub enabled: bool,
    /// Whether circuit breaker is currently open
    pub is_open: bool,
    /// Current failure count
    pub failure_count: u32,
    /// When the circuit breaker was opened
    pub opened_at: Option<Instant>,
}

/// Get the global multilevel cache instance
pub async fn get_client() -> Arc<MultilevelClient> {
    MULTILEVEL_CACHE
        .get()
        .expect("Multilevel cache not initialized")
        .clone()
}

/// Initialize the multilevel cache with configuration
pub async fn init(
    local: Arc<LocalClient>,
    redis: Arc<RedisClient>,
    config: MultilevelConfig,
) -> Result<Arc<MultilevelClient>> {
    let client = MultilevelClient::new(local, redis, config);
    let client = Arc::new(client);

    if MULTILEVEL_CACHE.set(client.clone()).is_err() {
        return Err(MultilevelError::AlreadyInitialized);
    }

    Ok(client)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{LocalConfig, PoolConfig, RedisConfig, RedisMode, RedisNode};

    async fn create_test_clients() -> (Arc<LocalClient>, Arc<RedisClient>) {
        let redis_config = RedisConfig {
            mode: RedisMode::Standalone,
            node: Some(RedisNode {
                host: "localhost".to_string(),
                port: 6379,
            }),
            sentinel: None,
            cluster: None,
            pool: PoolConfig::default(),
            password: Some("1qaz!QAZ".to_string()),
            database: None,
        };

        let local_config = LocalConfig {
            max_capacity: 100,
            ttl: 60,
            tti: 30,
        };

        let redis_client = Arc::new(RedisClient::new(redis_config).await.unwrap());
        let local_client = Arc::new(LocalClient::new(local_config));

        (local_client, redis_client)
    }

    #[tokio::test]
    async fn test_local_mode() {
        let (local_client, redis_client) = create_test_clients().await;

        let config = MultilevelConfig {
            mode: CacheMode::Local,
            write_through: true,
            read_through: true,
            enable_circuit_breaker: true,
            failure_threshold: 5,
            failure_window_secs: 60,
            reset_timeout_secs: 300,
        };

        let client = MultilevelClient::new(local_client, redis_client, config);

        // 测试本地模式下的基本操作
        client.set("local_key", "local_value").await.unwrap();
        let value = client.get("local_key").await.unwrap();
        assert_eq!(value, Some("local_value".to_string()));

        // 验证数据只存在于本地缓存
        let redis_value = client.redis.get("local_key").await.unwrap();
        assert_eq!(redis_value, None);

        // 测试删除操作
        let deleted = client.del("local_key").await.unwrap();
        assert!(deleted);
        let value = client.get("local_key").await.unwrap();
        assert_eq!(value, None);
    }

    #[tokio::test]
    async fn test_redis_mode() {
        let (local_client, redis_client) = create_test_clients().await;

        let config = MultilevelConfig {
            mode: CacheMode::Redis,
            write_through: true,
            read_through: true,
            enable_circuit_breaker: true,
            failure_threshold: 5,
            failure_window_secs: 60,
            reset_timeout_secs: 300,
        };

        let client = MultilevelClient::new(local_client, redis_client, config);

        // 测试Redis模式下的基本操作
        client.set("redis_key", "redis_value").await.unwrap();
        let value = client.get("redis_key").await.unwrap();
        assert_eq!(value, Some("redis_value".to_string()));

        // 验证数据只存在于Redis缓存
        let local_value = client.local.get("redis_key").await.unwrap();
        assert_eq!(local_value, None);

        // 测试删除操作
        let deleted = client.del("redis_key").await.unwrap();
        assert!(deleted);
        let value = client.get("redis_key").await.unwrap();
        assert_eq!(value, None);
    }

    #[tokio::test]
    async fn test_multi_mode() {
        let (local_client, redis_client) = create_test_clients().await;

        let config = MultilevelConfig {
            mode: CacheMode::Multi,
            write_through: true,
            read_through: true,
            enable_circuit_breaker: true,
            failure_threshold: 5,
            failure_window_secs: 60,
            reset_timeout_secs: 300,
        };

        let client = MultilevelClient::new(local_client, redis_client, config);

        // 测试写入，应同时写入本地和Redis
        client.set("multi_key", "multi_value").await.unwrap();

        // 验证两级缓存都有数据
        let local_value = client.local.get("multi_key").await.unwrap();
        let redis_value = client.redis.get("multi_key").await.unwrap();
        assert_eq!(local_value, Some("multi_value".to_string()));
        assert_eq!(redis_value, Some("multi_value".to_string()));

        // 测试本地缓存命中
        client.local.clear().await;
        client.set("local_hit", "hit_value").await.unwrap();
        let value = client.get("local_hit").await.unwrap();
        assert_eq!(value, Some("hit_value".to_string()));

        // 测试Redis回填本地缓存
        client.local.clear().await;
        let value = client.get("multi_key").await.unwrap();
        assert_eq!(value, Some("multi_value".to_string()));
        // 验证数据已回填到本地缓存
        let local_value = client.local.get("multi_key").await.unwrap();
        assert_eq!(local_value, Some("multi_value".to_string()));
    }

    #[tokio::test]
    async fn test_stats_and_clear() {
        let (local_client, redis_client) = create_test_clients().await;

        let config = MultilevelConfig {
            mode: CacheMode::Multi,
            write_through: true,
            read_through: true,
            enable_circuit_breaker: true,
            failure_threshold: 5,
            failure_window_secs: 60,
            reset_timeout_secs: 300,
        };

        let client = MultilevelClient::new(local_client, redis_client, config);

        // 写入一些测试数据并立即验证
        for i in 0..5 {
            let key = format!("key_{}", i);
            let value = format!("value_{}", i);

            // 写入数据
            client.set(&key, &value).await.unwrap();

            // 立即验证数据已写入本地缓存
            let local_value = client.local.get(&key).await.unwrap();
            assert_eq!(
                local_value,
                Some(value.clone()),
                "Data should be in local cache immediately after set: key={}",
                key
            );

            // 读取数据以产生命中统计
            let _ = client.get(&key).await.unwrap();
            let _ = client.get(&key).await.unwrap();
        }
        // 运行维护任务
        client.local.run_pending_tasks().await;

        // 验证统计信息
        let local_stats = client.local.stats();
        assert_eq!(local_stats.hits, 15);
        assert_eq!(local_stats.misses, 0);
        assert_eq!(local_stats.size, 5);

        // 测试清除功能
        client.clear().await.unwrap();

        // 验证数据已被清除
        for i in 0..5 {
            let key = format!("key_{}", i);
            let value = client.get(&key).await.unwrap();
            assert_eq!(
                value, None,
                "Cache should be empty after clear: key={}",
                key
            );
        }

        // 验证清除后的缓存大小
        let final_stats = client.local.stats();
        println!("final_stats: {:?}", final_stats);
        assert_eq!(final_stats.hits, 15);
        assert_eq!(final_stats.misses, 5);
        assert_eq!(final_stats.size, 0);
    }

    #[tokio::test]
    async fn test_circuit_breaker() {
        let (local_client, redis_client) = create_test_clients().await;

        let config = MultilevelConfig {
            mode: CacheMode::Multi,
            write_through: true,
            read_through: true,
            enable_circuit_breaker: true,
            failure_threshold: 2, // 设置较小的阈值便于测试
            failure_window_secs: 60,
            reset_timeout_secs: 1, // 设置较小的重置时间便于测试
        };

        let client =
            MultilevelClient::new(local_client.clone(), redis_client.clone(), config.clone());

        // 首先正常写入一些数据
        client.set("cb_key", "cb_value").await.unwrap();

        // 模拟Redis故障（通过修改密码使Redis连接失败）
        let bad_redis_config = RedisConfig {
            mode: RedisMode::Standalone,
            node: Some(RedisNode {
                host: "localhost".to_string(),
                port: 6379,
            }),
            sentinel: None,
            cluster: None,
            pool: PoolConfig::default(),
            password: Some("wrong_password".to_string()),
            database: None,
        };
        let bad_redis_client = Arc::new(RedisClient::new(bad_redis_config).await.unwrap());
        let client = MultilevelClient::new(local_client, bad_redis_client, config);

        // 触发熔断器
        for i in 0..3 {
            let result = client.set("trigger_key", "trigger_value").await;
            assert!(
                result.is_ok(),
                "Set operation should succeed with local fallback"
            );

            // 检查熔断器状态
            let status = client.circuit_breaker_status().await;
            log::info!(
                "Circuit breaker status after failure {}: open={}, count={}, opened_at={:?}",
                i + 1,
                status.is_open,
                status.failure_count,
                status.opened_at
            );
        }

        // 验证熔断器状态
        let status = client.circuit_breaker_status().await;
        assert!(status.is_open, "Expected circuit breaker to be open");
        assert!(status.failure_count >= 2, "Expected at least 2 failures");
        log::info!(
            "Circuit breaker status before reset: open={}, count={}, opened_at={:?}",
            status.is_open,
            status.failure_count,
            status.opened_at
        );

        // 验证在熔断状态下仍然可以使用本地缓存
        client.set("local_key", "local_value").await.unwrap();
        let value = client.get("local_key").await.unwrap();
        assert_eq!(value, Some("local_value".to_string()));

        // 运行维护任务
        client.local.run_pending_tasks().await;

        // 验证熔断器已重置
        let status = client.circuit_breaker_status().await;
        log::info!(
            "Circuit breaker status after reset: open={}, count={}, opened_at={:?}",
            status.is_open,
            status.failure_count,
            status.opened_at
        );
        // 等待重置时间
        tokio::time::sleep(Duration::from_secs(2)).await;

        // 尝试一次 Redis 操作，触发重置检查
        let _ = client.is_redis_available().await;

        // 再次检查熔断器状态
        let final_status = client.circuit_breaker_status().await;
        log::info!(
            "Circuit breaker final status: open={}, count={}, opened_at={:?}",
            final_status.is_open,
            final_status.failure_count,
            final_status.opened_at
        );

        assert!(
            !final_status.is_open,
            "Expected circuit breaker to be closed after reset timeout"
        );
        assert_eq!(
            final_status.failure_count, 0,
            "Expected failure count to be reset to 0"
        );

        // 验证可以继续写入本地缓存
        client.set("after_reset", "value").await.unwrap();
        let value = client.get("after_reset").await.unwrap();
        assert_eq!(value, Some("value".to_string()));
    }

    #[tokio::test]
    async fn test_write_through_and_read_through() {
        let (local_client, redis_client) = create_test_clients().await;

        let config = MultilevelConfig {
            mode: CacheMode::Multi,
            write_through: false, // 禁用写透
            read_through: false,  // 禁用读透
            enable_circuit_breaker: true,
            failure_threshold: 5,
            failure_window_secs: 60,
            reset_timeout_secs: 300,
        };

        let client = MultilevelClient::new(local_client, redis_client, config);

        // 测试禁用写透时的行为
        client.set("wt_key", "wt_value").await.unwrap();

        // Redis应该有数据，但本地缓存没有
        let redis_value = client.redis.get("wt_key").await.unwrap();
        let local_value = client.local.get("wt_key").await.unwrap();
        assert_eq!(redis_value, Some("wt_value".to_string()));
        assert_eq!(local_value, None);

        // 测试禁用读透时的行为
        let value = client.get("wt_key").await.unwrap();
        assert_eq!(value, Some("wt_value".to_string()));

        // 本地缓存应该仍然为空
        let local_value = client.local.get("wt_key").await.unwrap();
        assert_eq!(local_value, None);
    }
}