cm-dashboard/agent/src/metrics/mod.rs

use anyhow::Result;
use cm_dashboard_shared::{Metric, StatusTracker};
use tracing::{error, info};

use crate::collectors::{
    backup::BackupCollector, cpu::CpuCollector, disk::DiskCollector, memory::MemoryCollector,
    nixos::NixOSCollector, systemd::SystemdCollector, Collector,
};
use crate::config::{AgentConfig, CollectorConfig};

/// Manages all metric collectors
pub struct MetricCollectionManager {
    collectors: Vec<Box<dyn Collector>>,
    status_tracker: StatusTracker,
}

impl MetricCollectionManager {
    pub async fn new(config: &CollectorConfig, _agent_config: &AgentConfig) -> Result<Self> {
        let mut collectors: Vec<Box<dyn Collector>> = Vec::new();

        // Benchmark mode - only enable specific collector based on env var
        let benchmark_mode = std::env::var("BENCHMARK_COLLECTOR").ok();

        match benchmark_mode.as_deref() {
            Some("cpu") => {
                // CPU collector only
                if config.cpu.enabled {
                    let cpu_collector = CpuCollector::new(config.cpu.clone());
                    collectors.push(Box::new(cpu_collector));
                    info!("BENCHMARK: CPU collector only");
                }
            }
            Some("memory") => {
                // Memory collector only
                if config.memory.enabled {
                    let memory_collector = MemoryCollector::new(config.memory.clone());
                    collectors.push(Box::new(memory_collector));
                    info!("BENCHMARK: Memory collector only");
                }
            }
            Some("disk") => {
                // Disk collector only
                let disk_collector = DiskCollector::new(config.disk.clone());
                collectors.push(Box::new(disk_collector));
                info!("BENCHMARK: Disk collector only");
            }
            Some("systemd") => {
                // Systemd collector only
                let systemd_collector = SystemdCollector::new(config.systemd.clone());
                collectors.push(Box::new(systemd_collector));
                info!("BENCHMARK: Systemd collector only");
            }
            Some("backup") => {
                // Backup collector only
                if config.backup.enabled {
                    let backup_collector = BackupCollector::new(
                        config.backup.backup_paths.first().cloned(),
                        config.backup.max_age_hours,
                    );
                    collectors.push(Box::new(backup_collector));
                    info!("BENCHMARK: Backup collector only");
                }
            }
            Some("none") => {
                // No collectors - test agent loop only
                info!("BENCHMARK: No collectors enabled");
            }
            _ => {
                // Normal mode - all collectors
                if config.cpu.enabled {
                    let cpu_collector = CpuCollector::new(config.cpu.clone());
                    collectors.push(Box::new(cpu_collector));
                    info!("CPU collector initialized");
                }

                if config.memory.enabled {
                    let memory_collector = MemoryCollector::new(config.memory.clone());
                    collectors.push(Box::new(memory_collector));
                    info!("Memory collector initialized");
                }

                let disk_collector = DiskCollector::new(config.disk.clone());
                collectors.push(Box::new(disk_collector));
                info!("Disk collector initialized");

                let systemd_collector = SystemdCollector::new(config.systemd.clone());
                collectors.push(Box::new(systemd_collector));
                info!("Systemd collector initialized");

                if config.backup.enabled {
                    let backup_collector = BackupCollector::new(
                        config.backup.backup_paths.first().cloned(),
                        config.backup.max_age_hours,
                    );
                    collectors.push(Box::new(backup_collector));
                    info!("Backup collector initialized");
                }

                if config.nixos.enabled {
                    let nixos_collector = NixOSCollector::new(config.nixos.clone());
                    collectors.push(Box::new(nixos_collector));
                    info!("NixOS collector initialized");
                }
            }
        }

        info!(
            "Metric collection manager initialized with {} collectors",
            collectors.len()
        );

        Ok(Self {
            collectors,
            status_tracker: StatusTracker::new(),
        })
    }

    /// Force collection from ALL collectors immediately (used at startup)
    pub async fn collect_all_metrics_force(&mut self) -> Result<Vec<Metric>> {
        self.collect_all_metrics().await
    }

    /// Collect metrics from all collectors
    pub async fn collect_all_metrics(&mut self) -> Result<Vec<Metric>> {
        let mut all_metrics = Vec::new();

        for collector in &self.collectors {
            match collector.collect(&mut self.status_tracker).await {
                Ok(metrics) => {
                    all_metrics.extend(metrics);
                }
                Err(e) => {
                    error!("Collector failed: {}", e);
                }
            }
        }
        Ok(all_metrics)
    }

}