mocra-core 0.4.0

use crate::cacheable::CacheService;
use crate::common::model::Request;
use crate::common::model::config::Config;
use crate::common::model::download_config::DownloadConfig;
use crate::downloader::request_downloader::RequestDownloader;
use crate::downloader::{Downloader, WebSocketDownloader};
use crate::utils::distributed_rate_limit::DistributedSlidingWindowRateLimiter;
use crate::utils::redis_lock::DistributedLockManager;
use dashmap::DashMap;
use deadpool_redis::redis::Script;
use std::sync::Arc;
use std::time::SystemTime;
use tokio::sync::RwLock;

pub struct DownloaderManager {
    /// 应用配置(命名空间名等);此前吃整个 `Arc<State>`,现窄化为具体依赖(重构 Phase 2)。
    pub app_config: Arc<RwLock<Config>>,
    /// 分布式锁管理器(用于取 Redis 连接池)。
    pub locker: Arc<DistributedLockManager>,
    /// 默认下载器(缺省 reqwest);可经 [`set_default_downloader`](Self::set_default_downloader)
    /// 在启动前替换(如换成浏览器渲染 / 代理轮换 / 自定义重试的下载器)。
    pub default_downloader: RwLock<Box<dyn Downloader>>,
    // Registered downloader factories.
    pub downloader: Arc<DashMap<String, Box<dyn Downloader>>>,
    // Task downloader configuration.
    pub config: Arc<DashMap<String, DownloadConfig>>,
    // Task downloader instances.
    pub task_downloader: Arc<DashMap<String, Box<dyn Downloader>>>,
    pub wss_downloader: Arc<WebSocketDownloader>,
    // Records the last expiration update timestamp to reduce Redis write frequency.
    pub expire_update_cache: Arc<DashMap<String, u64>>,
}

impl DownloaderManager {
    /// Create a new DownloaderManager instance.
    ///
    /// # Arguments
    /// * `state` - Shared application state containing configuration, rate limiter, etc.
    pub async fn new(
        app_config: Arc<RwLock<Config>>,
        limiter: Arc<DistributedSlidingWindowRateLimiter>,
        locker: Arc<DistributedLockManager>,
        cache_service: Arc<CacheService>,
    ) -> Self {
        let (pool_size, max_response_size) = {
            let cfg = app_config.read().await;
            (
                cfg.download_config.pool_size.unwrap_or(200),
                cfg.download_config
                    .max_response_size
                    .unwrap_or(10 * 1024 * 1024),
            )
        };

        DownloaderManager {
            app_config,
            locker: locker.clone(),
            default_downloader: RwLock::new(Box::new(RequestDownloader::new(
                Arc::clone(&limiter),
                Arc::clone(&locker),
                Arc::clone(&cache_service),
                pool_size,
                max_response_size,
            ))),
            // Downloader factory list.
            downloader: Arc::new(DashMap::new()),
            // Task downloader configuration.
            config: Arc::new(DashMap::new()),
            // Task downloader instances.
            task_downloader: Arc::new(DashMap::new()),
            wss_downloader: Arc::new(WebSocketDownloader::new()),
            expire_update_cache: Arc::new(DashMap::new()),
        }
    }

    /// Register a custom downloader implementation.
    ///
    /// The downloader is selected for a module/request when its `DownloadConfig.downloader`
    /// name matches this downloader's [`name()`](Downloader::name).
    pub async fn register(&self, downloader: Box<dyn Downloader>) {
        self.downloader.insert(downloader.name(), downloader);
    }

    /// Replaces the default downloader (used when a request's `config.downloader` does not
    /// match any registered downloader). Set this before the engine starts.
    pub async fn set_default_downloader(&self, downloader: Box<dyn Downloader>) {
        *self.default_downloader.write().await = downloader;
    }

    /// Set rate limit for a specific limit_id dynamically.
    ///
    /// This updates the rate limit configuration for an active downloader instance.
    pub async fn set_limit(&self, limit_id: &str, limit: f32) {
        let downloader = self
            .task_downloader
            .get(limit_id)
            .map(|d| dyn_clone::clone_box(d.value().as_ref()));
        // guard dropped here — never hold DashMap Ref across .await
        if let Some(d) = downloader {
            d.set_limit(limit_id, limit).await;
        }
    }

    // Helper function to get the current timestamp.
    fn current_timestamp() -> u64 {
        SystemTime::now()
            .duration_since(SystemTime::UNIX_EPOCH)
            .unwrap()
            .as_secs()
    }

    // Cleans up expired downloaders.
    async fn cleanup_expired_downloader(&self) {
        let downloader_expire_time = self
            .app_config
            .read()
            .await
            .download_config
            .downloader_expire;

        let current_time = Self::current_timestamp();
        let max_score = current_time.saturating_sub(downloader_expire_time);

        let config_name = self.app_config.read().await.name.clone();
        let key = format!("{}:downloader_expire", config_name);

        // Lua script to get and remove expired keys
        let script = Script::new(
            r#"
            local key = KEYS[1]
            local max_score = ARGV[1]
            local expired = redis.call('ZRANGEBYSCORE', key, '-inf', max_score)
            if #expired > 0 then
                redis.call('ZREM', key, unpack(expired))
            end
            return expired
        "#,
        );

        let mut expired_keys: Vec<String> = Vec::new();

        // Execute Lua script
        let pool = self.locker.get_pool();
        if let Some(pool) = pool
            && let Ok(mut conn) = pool.get().await
        {
            let result: Result<Vec<String>, _> = script
                .key(&key)
                .arg(max_score)
                .invoke_async(&mut conn)
                .await;

            if let Ok(keys) = result {
                expired_keys = keys;
            }
        }

        // Remove expired keys from local map
        for key in &expired_keys {
            self.task_downloader.remove(key);
        }

        // Check health status.
        let check_list: Vec<(String, Box<dyn Downloader>)> = self
            .task_downloader
            .iter()
            .map(|r| (r.key().clone(), dyn_clone::clone_box(r.value().as_ref())))
            .collect();

        for (key, downloader) in check_list {
            if downloader.health_check().await.is_err() {
                self.task_downloader.remove(&key);
                if let Some(pool) = pool
                    && let Ok(mut conn) = pool.get().await
                {
                    let _: () = deadpool_redis::redis::cmd("ZREM")
                        .arg(&key)
                        .arg(&key)
                        .query_async(&mut conn)
                        .await
                        .unwrap_or(());
                }
            }
        }
    }

    /// Start the background cleanup task.
    /// Should be called once at startup.
    pub fn start_background_cleaner(self: Arc<Self>) {
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(std::time::Duration::from_secs(60));
            loop {
                interval.tick().await;
                self.cleanup_expired_downloader().await;
            }
        });
    }

    /// Get or create a downloader for the given request.
    ///
    /// This method manages the lifecycle of downloader instances for specific tasks (modules).
    /// It handles creation, configuration updates, and expiration of idle downloaders.
    ///
    /// # Arguments
    /// * `request` - The request containing module ID and limit ID.
    /// * `download_config` - Configuration for the downloader.
    pub async fn get_downloader(
        &self,
        request: &Request,
        download_config: DownloadConfig,
    ) -> Box<dyn Downloader> {
        let current_time = Self::current_timestamp();
        let module_id = request.module_id();

        // Check whether expiration time needs to be updated (once every 60 seconds).
        let should_update = if let Some(last_update) = self.expire_update_cache.get(&module_id) {
            current_time > *last_update + 60
        } else {
            true
        };

        if should_update {
            // Optimistically update cache to prevent spamming the spawn
            self.expire_update_cache
                .insert(module_id.clone(), current_time);

            // Update expiration time (Redis ZSET).
            let config_name = self.app_config.read().await.name.clone();
            let key = format!("{}:downloader_expire", config_name);
            let pool = self.locker.get_pool().cloned();
            let module_id_clone = module_id.clone();
            let current_time_clone = current_time;

            tokio::spawn(async move {
                if let Some(pool) = pool
                    && let Ok(mut conn) = pool.get().await
                {
                    let _: () = deadpool_redis::redis::cmd("ZADD")
                        .arg(&key)
                        .arg(current_time_clone)
                        .arg(&module_id_clone)
                        .query_async(&mut conn)
                        .await
                        .unwrap_or(());
                }
            });
        }

        // Get or insert configuration.
        // Extract the cached value and drop the DashMap guard BEFORE any
        // potential .insert() on the same map — holding a Ref (read guard)
        // while calling .insert() (write lock) on the same shard is a
        // parking_lot self-deadlock.
        let cached_config = self.config.get(&module_id).map(|existing| existing.clone());
        // guard dropped here
        let config = if let Some(cached) = cached_config {
            if cached != download_config {
                self.config
                    .insert(module_id.clone(), download_config.clone());
                download_config.clone()
            } else {
                cached
            }
        } else {
            self.config
                .insert(module_id.clone(), download_config.clone());
            download_config.clone()
        };

        // Determine effective limit_id here to ensure rate limiter gets correct key
        let limit_id = if request.limit_id.is_empty() {
            request.module_id()
        } else {
            request.limit_id.clone()
        };

        // Get or create downloader.
        // Clone the cached downloader and drop the DashMap guard before any .await.
        // Holding a DashMap Ref across an await can deadlock the Tokio runtime
        // (parking_lot RwLock blocks the OS thread).
        let cached_downloader = self
            .task_downloader
            .get(&module_id)
            .map(|d| dyn_clone::clone_box(d.value().as_ref()));
        // guard dropped here
        if let Some(d) = cached_downloader {
            d.set_config(&limit_id, config).await;
            return d;
        }

        let new_downloader = if let Some(registered) = self.downloader.get(&config.downloader) {
            dyn_clone::clone_box(registered.value().as_ref())
        } else {
            dyn_clone::clone_box(self.default_downloader.read().await.as_ref())
        };

        // Ensure the configuration is up to date.
        new_downloader.set_config(&limit_id, config).await;
        self.task_downloader.insert(
            module_id.clone(),
            dyn_clone::clone_box(new_downloader.as_ref()),
        );
        new_downloader
    }
    /// Clear all configurations and downloaders.
    ///
    /// This removes all registered task downloaders and their configurations.
    pub async fn clear(&self) {
        self.config.clear();
        self.task_downloader.clear();
    }
}