Mantis/mantis/src/core/infrastructure/health.rs

use std::sync::Arc;
use std::time::Duration;

use sysinfo::{Components, Networks, System};
use tokio::sync::{broadcast, oneshot, RwLock};
use tokio::time::interval;
use macros::log;

use crate::core::infrastructure::app_config::AppConfig;
use crate::model::log::health::Health;
use crate::model::error::Error;
use crate::model::health::{
    ConfiguredNetworkStats,
    CpuCoreInfo,
    CpuDetails,
    LoadAverage,
    MemoryUsage,
    NetworkStats,
    SystemHealthMetrics,
    SystemHealthStatus,
    SystemInfo
};

pub struct SystemHealth {
    system: RwLock<System>,
    networks: RwLock<Networks>,
    components: RwLock<Components>,
    broadcast_tx: broadcast::Sender<SystemHealthMetrics>,
    ingress_interface: String,
    egress_interface: String,
}


impl SystemHealth {
    pub fn new(config: Arc<AppConfig>) -> Result<Self, Error> {
        let (broadcast_tx, _) = broadcast::channel(100);

        let health = SystemHealth {
            system: RwLock::new(System::new_all()),
            networks: RwLock::new(Networks::new_with_refreshed_list()),
            components: RwLock::new(Components::new_with_refreshed_list()),
            broadcast_tx,
            ingress_interface: config.ingress_ifname.clone(),
            egress_interface: config.egress_ifname.clone(),
        };

        Ok(health)
    }

    pub async fn run(self: Arc<Self>, monitoring_interval: Duration) -> oneshot::Sender<()> {
        let (sender, mut receiver) = oneshot::channel();
        let health = self.clone();

        tokio::spawn(async move {
            let mut interval_timer = interval(monitoring_interval);

            loop {
                tokio::select! {
                    biased;
                    _ = &mut receiver => {
                        break;
                    }
                    _ = interval_timer.tick() => {
                        health.refresh_and_broadcast().await;
                    }
                }
            }
        });

        sender
    }

    async fn refresh_and_broadcast(&self) {
        self.system.write().await.refresh_all();
        self.networks.write().await.refresh(true);
        self.components.write().await.refresh(true);

        let system = self.system.read().await;
        let networks = self.networks.read().await;
        let components = self.components.read().await;

        let metrics = Self::collect_metrics(
            &system,
            &networks,
            &components,
            &self.ingress_interface,
            &self.egress_interface,
        );

        drop(system);
        drop(networks);
        drop(components);

        if self.broadcast_tx.receiver_count() > 0 {
            if let Err(e) = self.broadcast_tx.send(metrics) {
                log!(Health::BroadcastFailed(e.to_string()));
            }
        }
    }

    fn collect_metrics(
        system: &System,
        networks: &Networks,
        components: &Components,
        ingress_interface: &str,
        egress_interface: &str,
    ) -> SystemHealthMetrics {
        let timestamp = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .expect("system time is after UNIX_EPOCH")
            .as_secs();

        let boot_time = System::boot_time();
        let uptime_seconds = timestamp - boot_time;

        let system_info = Self::collect_system_info(system);
        let cpu_details = Self::collect_cpu_details(system);

        let memory_usage = MemoryUsage {
            total: system.total_memory(),
            used: system.used_memory(),
            available: system.available_memory(),
            usage_percent: (system.used_memory() as f32 / system.total_memory() as f32) * 100.0,
            swap_total: system.total_swap(),
            swap_used: system.used_swap(),
        };

        let network_stats = Self::collect_configured_network_stats(
            networks,
            ingress_interface,
            egress_interface,
        );

        let load_average = System::load_average();
        let load_average = if load_average.one != 0.0 || load_average.five != 0.0 || load_average.fifteen != 0.0 {
            Some(LoadAverage {
                one_minute: load_average.one,
                five_minute: load_average.five,
                fifteen_minute: load_average.fifteen,
            })
        } else {
            None
        };

        let temperature = components
            .iter()
            .find(|component| {
                let label = component.label().to_lowercase();
                label.contains("cpu") || label.contains("core") || label.contains("processor")
            })
            .and_then(|component| component.temperature());

        SystemHealthMetrics {
            timestamp,
            boot_time,
            uptime_seconds,
            system_info,
            cpu_details,
            memory_usage,
            network_stats,
            load_average,
            temperature,
        }
    }

    fn collect_system_info(system: &System) -> SystemInfo {
        SystemInfo {
            kernel_version: System::kernel_version(),
            os_name: System::name(),
            os_version: System::os_version(),
            architecture: std::env::consts::ARCH.to_string(),
            total_processes: system.processes().len(),
        }
    }

    fn collect_cpu_details(system: &System) -> CpuDetails {
        let cpus = system.cpus();

        let cpu_usage = cpus.iter().map(|cpu| cpu.cpu_usage()).sum::<f32>() / cpus.len() as f32;

        let cores: Vec<CpuCoreInfo> = cpus
            .iter()
            .enumerate()
            .map(|(index, cpu)| CpuCoreInfo {
                core_id: index,
                usage_percent: cpu.cpu_usage(),
                frequency: cpu.frequency(),
            })
            .collect();

        let cpu_brand = cpus
            .first()
            .map(|cpu| cpu.brand().to_string())
            .unwrap_or_else(|| "Unknown".to_string());

        let avg_frequency = if !cores.is_empty() {
            cores.iter().map(|core| core.frequency).sum::<u64>() / cores.len() as u64
        } else {
            0
        };

        CpuDetails {
            cpu_brand,
            core_count: cores.len(),
            cpu_usage,
            cpu_frequency: avg_frequency,
            cores,
        }
    }

    fn collect_configured_network_stats(
        networks: &Networks,
        ingress_interface: &str,
        egress_interface: &str,
    ) -> ConfiguredNetworkStats {
        let create_network_stats = |interface_name: &str| -> Option<NetworkStats> {
            networks.get(interface_name).map(|network| NetworkStats {
                interface: interface_name.to_string(),
                bytes_received: network.total_received(),
                bytes_transmitted: network.total_transmitted(),
                packets_received: network.total_packets_received(),
                packets_transmitted: network.total_packets_transmitted(),
                errors_received: network.total_errors_on_received(),
                errors_transmitted: network.total_errors_on_transmitted(),
            })
        };

        let ingress = create_network_stats(ingress_interface);
        let egress = create_network_stats(egress_interface);

        if ingress.is_none() {
            log!(Health::InterfaceNotFound("Ingress".to_string(), ingress_interface.to_string()));
        }
        if egress.is_none() {
            log!(Health::InterfaceNotFound("Egress".to_string(), egress_interface.to_string()));
        }

        ConfiguredNetworkStats {
            ingress,
            egress,
        }
    }

    pub async fn get_current_metrics(&self) -> SystemHealthMetrics {
        self.system.write().await.refresh_all();
        self.networks.write().await.refresh(true);
        self.components.write().await.refresh(true);

        let system = self.system.read().await;
        let networks = self.networks.read().await;
        let components = self.components.read().await;

        Self::collect_metrics(
            &system,
            &networks,
            &components,
            &self.ingress_interface,
            &self.egress_interface,
        )
    }

    pub fn subscribe_to_metrics(&self) -> broadcast::Receiver<SystemHealthMetrics> {
        self.broadcast_tx.subscribe()
    }

    pub async fn is_system_healthy(&self) -> SystemHealthStatus {
        let metrics = self.get_current_metrics().await;

        let mut status = SystemHealthStatus {
            overall_healthy: true,
            issues: Vec::new(),
            warnings: Vec::new(),
        };

        if metrics.cpu_details.cpu_usage > 90.0 {
            status.overall_healthy = false;
            status
                .issues
                .push(format!("High CPU usage: {:.1}%", metrics.cpu_details.cpu_usage));
        } else if metrics.cpu_details.cpu_usage > 75.0 {
            status
                .warnings
                .push(format!("Moderate CPU usage: {:.1}%", metrics.cpu_details.cpu_usage));
        }

        if metrics.memory_usage.usage_percent > 95.0 {
            status.overall_healthy = false;
            status.issues.push(format!(
                "Critical memory usage: {:.1}%",
                metrics.memory_usage.usage_percent
            ));
        } else if metrics.memory_usage.usage_percent > 80.0 {
            status.warnings.push(format!(
                "High memory usage: {:.1}%",
                metrics.memory_usage.usage_percent
            ));
        }

        if let Some(temp) = metrics.temperature {
            if temp > 80.0 {
                status.overall_healthy = false;
                status
                    .issues
                    .push(format!("High CPU temperature: {:.1}°C", temp));
            } else if temp > 70.0 {
                status
                    .warnings
                    .push(format!("Elevated CPU temperature: {:.1}°C", temp));
            }
        }

        if metrics.network_stats.ingress.is_none() {
            status.overall_healthy = false;
            status.issues.push("Ingress interface not available".to_string());
        }
        if metrics.network_stats.egress.is_none() {
            status.overall_healthy = false;
            status.issues.push("Egress interface not available".to_string());
        }

        status
    }
}