mocra-core 0.4.1

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::lock::DistributedLockManager;
use dashmap::DashMap;
use std::sync::Arc;
use std::time::SystemTime;
use tokio::sync::RwLock;

pub struct DownloaderManager {
    /// Application config (namespace name, etc.); it used to take the entire `Arc<State>`, now
    /// narrowed to the concrete dependency (refactor Phase 2).
    pub app_config: Arc<RwLock<Config>>,
    /// Distributed lock manager (coordination backend / in-process local locks).
    pub locker: Arc<DistributedLockManager>,
    /// Default downloader (reqwest by default); can be swapped out before startup via
    /// [`set_default_downloader`](Self::set_default_downloader) (e.g. for a downloader that does
    /// browser rendering / proxy rotation / custom retries).
    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 throttle expiry index writes.
    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);

        // Evict expired downloaders locally by last-used time (the in-process
        // `expire_update_cache` acts as the index).
        let mut expired_keys: Vec<String> = Vec::new();
        for entry in self.expire_update_cache.iter() {
            if *entry.value() < max_score {
                expired_keys.push(entry.key().clone());
            }
        }
        for key in &expired_keys {
            self.task_downloader.remove(key);
            self.expire_update_cache.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);
                self.expire_update_cache.remove(&key);
            }
        }
    }

    /// 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 {
            // Record last-use time (drives local expiry in cleanup_expired_downloader).
            self.expire_update_cache
                .insert(module_id.clone(), current_time);
        }

        // 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();
    }
}