cm-dashboard/agent/src/collectors/system.rs

use async_trait::async_trait;
use serde_json::json;
use std::time::Duration;
use tokio::fs;
use tokio::process::Command;
use tracing::debug;

use super::{Collector, CollectorError, CollectorOutput, AgentType};

pub struct SystemCollector {
    enabled: bool,
    interval: Duration,
}

impl SystemCollector {
    pub fn new(enabled: bool, interval_ms: u64) -> Self {
        Self {
            enabled,
            interval: Duration::from_millis(interval_ms),
        }
    }

    async fn get_cpu_load(&self) -> Result<(f32, f32, f32), CollectorError> {
        let output = Command::new("/run/current-system/sw/bin/uptime")
            .output()
            .await
            .map_err(|e| CollectorError::CommandFailed { 
                command: "uptime".to_string(), 
                message: e.to_string() 
            })?;

        let uptime_str = String::from_utf8_lossy(&output.stdout);
        
        // Parse load averages from uptime output
        // Format with comma decimals: "... load average: 3,30, 3,17, 2,84"
        if let Some(load_part) = uptime_str.split("load average:").nth(1) {
            // Use regex or careful parsing for comma decimal separator locale
            let load_str = load_part.trim();
            // Split on ", " to separate the three load values
            let loads: Vec<&str> = load_str.split(", ").collect();
            if loads.len() >= 3 {
                let load_1 = loads[0].trim().replace(',', ".").parse::<f32>()
                    .map_err(|_| CollectorError::ParseError { message: "Failed to parse 1min load".to_string() })?;
                let load_5 = loads[1].trim().replace(',', ".").parse::<f32>()
                    .map_err(|_| CollectorError::ParseError { message: "Failed to parse 5min load".to_string() })?;
                let load_15 = loads[2].trim().replace(',', ".").parse::<f32>()
                    .map_err(|_| CollectorError::ParseError { message: "Failed to parse 15min load".to_string() })?;
                
                return Ok((load_1, load_5, load_15));
            }
        }
        
        Err(CollectorError::ParseError { message: "Failed to parse load averages".to_string() })
    }

    async fn get_cpu_temperature(&self) -> Option<f32> {
        // Try to find CPU-specific thermal zones first (x86_pkg_temp, coretemp, etc.)
        for i in 0..10 {
            let type_path = format!("/sys/class/thermal/thermal_zone{}/type", i);
            let temp_path = format!("/sys/class/thermal/thermal_zone{}/temp", i);
            
            if let (Ok(zone_type), Ok(temp_str)) = (
                fs::read_to_string(&type_path).await,
                fs::read_to_string(&temp_path).await,
            ) {
                let zone_type = zone_type.trim();
                if let Ok(temp_millic) = temp_str.trim().parse::<f32>() {
                    let temp_c = temp_millic / 1000.0;
                    // Look for reasonable temperatures first
                    if temp_c > 20.0 && temp_c < 150.0 {
                        // Prefer CPU package temperature zones
                        if zone_type == "x86_pkg_temp" || zone_type.contains("coretemp") {
                            debug!("Found CPU temperature: {}°C from {} ({})", temp_c, temp_path, zone_type);
                            return Some(temp_c);
                        }
                    }
                }
            }
        }
        
        // Fallback: try any reasonable temperature if no CPU-specific zone found
        for i in 0..10 {
            let temp_path = format!("/sys/class/thermal/thermal_zone{}/temp", i);
            if let Ok(temp_str) = fs::read_to_string(&temp_path).await {
                if let Ok(temp_millic) = temp_str.trim().parse::<f32>() {
                    let temp_c = temp_millic / 1000.0;
                    if temp_c > 20.0 && temp_c < 150.0 {
                        debug!("Found fallback temperature: {}°C from {}", temp_c, temp_path);
                        return Some(temp_c);
                    }
                }
            }
        }
        None
    }

    async fn get_memory_info(&self) -> Result<(f32, f32), CollectorError> {
        let meminfo = fs::read_to_string("/proc/meminfo")
            .await
            .map_err(|e| CollectorError::IoError { message: format!("Failed to read /proc/meminfo: {}", e) })?;

        let mut total_kb = 0;
        let mut available_kb = 0;

        for line in meminfo.lines() {
            if line.starts_with("MemTotal:") {
                if let Some(value) = line.split_whitespace().nth(1) {
                    total_kb = value.parse::<u64>().unwrap_or(0);
                }
            } else if line.starts_with("MemAvailable:") {
                if let Some(value) = line.split_whitespace().nth(1) {
                    available_kb = value.parse::<u64>().unwrap_or(0);
                }
            }
        }

        if total_kb == 0 {
            return Err(CollectorError::ParseError { message: "Could not parse total memory".to_string() });
        }

        let total_mb = total_kb as f32 / 1024.0;
        let used_mb = total_mb - (available_kb as f32 / 1024.0);

        Ok((used_mb, total_mb))
    }

    async fn get_logged_in_users(&self) -> Option<Vec<String>> {
        // Get currently logged-in users using 'who' command
        let output = Command::new("who")
            .output()
            .await
            .ok()?;

        let who_output = String::from_utf8_lossy(&output.stdout);
        let mut users = Vec::new();
        
        for line in who_output.lines() {
            if let Some(username) = line.split_whitespace().next() {
                if !username.is_empty() && !users.contains(&username.to_string()) {
                    users.push(username.to_string());
                }
            }
        }
        
        if users.is_empty() {
            None
        } else {
            users.sort();
            Some(users)
        }
    }

    async fn get_cpu_cstate_info(&self) -> Option<Vec<String>> {
        // Read C-state information to show all sleep state distributions
        let mut cstate_times: Vec<(String, u64)> = Vec::new();
        let mut total_time = 0u64;
        
        // Check if C-state information is available
        if let Ok(mut entries) = fs::read_dir("/sys/devices/system/cpu/cpu0/cpuidle").await {
            while let Ok(Some(entry)) = entries.next_entry().await {
                let state_path = entry.path();
                let name_path = state_path.join("name");
                let time_path = state_path.join("time");
                
                if let (Ok(name), Ok(time_str)) = (
                    fs::read_to_string(&name_path).await,
                    fs::read_to_string(&time_path).await
                ) {
                    let name = name.trim().to_string();
                    if let Ok(time) = time_str.trim().parse::<u64>() {
                        total_time += time;
                        cstate_times.push((name, time));
                    }
                }
            }
            
            if total_time > 0 && !cstate_times.is_empty() {
                // Sort by C-state order: POLL, C1, C1E, C3, C6, C7s, C8, C9, C10
                cstate_times.sort_by(|a, b| {
                    let order_a = match a.0.as_str() {
                        "POLL" => 0,
                        "C1" => 1,
                        "C1E" => 2,
                        "C3" => 3,
                        "C6" => 4,
                        "C7s" => 5,
                        "C8" => 6,
                        "C9" => 7,
                        "C10" => 8,
                        _ => 99,
                    };
                    let order_b = match b.0.as_str() {
                        "POLL" => 0,
                        "C1" => 1,
                        "C1E" => 2,
                        "C3" => 3,
                        "C6" => 4,
                        "C7s" => 5,
                        "C8" => 6,
                        "C9" => 7,
                        "C10" => 8,
                        _ => 99,
                    };
                    order_a.cmp(&order_b)
                });
                
                // Format C-states as description lines (2 C-states per row)
                let mut result = Vec::new();
                let mut current_line = Vec::new();
                
                for (name, time) in cstate_times {
                    let percent = (time as f32 / total_time as f32) * 100.0;
                    if percent >= 0.1 { // Only show states with at least 0.1% time
                        current_line.push(format!("{}: {:.1}%", name, percent));
                        
                        // Split into rows when we have 2 items
                        if current_line.len() == 2 {
                            result.push(current_line.join(", "));
                            current_line.clear();
                        }
                    }
                }
                
                // Add remaining items as final line
                if !current_line.is_empty() {
                    result.push(current_line.join(", "));
                }
                
                return Some(result);
            }
        }
        
        None
    }

    fn determine_cpu_status(&self, cpu_load_5: f32) -> String {
        if cpu_load_5 >= 10.0 {
            "critical".to_string()
        } else if cpu_load_5 >= 9.0 {
            "warning".to_string()
        } else {
            "ok".to_string()
        }
    }

    fn determine_cpu_temp_status(&self, temp_c: f32) -> String {
        if temp_c >= 100.0 {
            "critical".to_string()
        } else if temp_c >= 100.0 {
            "warning".to_string()
        } else {
            "ok".to_string()
        }
    }

    fn determine_memory_status(&self, usage_percent: f32) -> String {
        if usage_percent >= 95.0 {
            "critical".to_string()
        } else if usage_percent >= 80.0 {
            "warning".to_string()
        } else {
            "ok".to_string()
        }
    }
}

#[async_trait]
impl Collector for SystemCollector {
    fn name(&self) -> &str {
        "system"
    }

    fn agent_type(&self) -> AgentType {
        AgentType::System
    }

    fn collect_interval(&self) -> Duration {
        self.interval
    }

    async fn collect(&self) -> Result<CollectorOutput, CollectorError> {
        if !self.enabled {
            return Err(CollectorError::ConfigError { message: "SystemCollector disabled".to_string() });
        }

        // Get CPU load averages
        let (cpu_load_1, cpu_load_5, cpu_load_15) = self.get_cpu_load().await?;
        let cpu_status = self.determine_cpu_status(cpu_load_5);

        // Get CPU temperature (optional) 
        let cpu_temp_c = self.get_cpu_temperature().await;
        let cpu_temp_status = cpu_temp_c.map(|temp| self.determine_cpu_temp_status(temp));

        // Get memory information
        let (memory_used_mb, memory_total_mb) = self.get_memory_info().await?;
        let memory_usage_percent = (memory_used_mb / memory_total_mb) * 100.0;
        let memory_status = self.determine_memory_status(memory_usage_percent);

        // Get C-state information (optional)
        let cpu_cstate_info = self.get_cpu_cstate_info().await;
        
        // Get logged-in users (optional)
        let logged_in_users = self.get_logged_in_users().await;

        let mut system_metrics = json!({
            "summary": {
                "cpu_load_1": cpu_load_1,
                "cpu_load_5": cpu_load_5,
                "cpu_load_15": cpu_load_15,
                "cpu_status": cpu_status,
                "memory_used_mb": memory_used_mb,
                "memory_total_mb": memory_total_mb,
                "memory_usage_percent": memory_usage_percent,
                "memory_status": memory_status,
            },
            "timestamp": chrono::Utc::now().timestamp() as u64,
        });

        // Add optional metrics if available
        if let Some(temp) = cpu_temp_c {
            system_metrics["summary"]["cpu_temp_c"] = json!(temp);
            if let Some(status) = cpu_temp_status {
                system_metrics["summary"]["cpu_temp_status"] = json!(status);
            }
        }

        if let Some(cstates) = cpu_cstate_info {
            system_metrics["summary"]["cpu_cstate"] = json!(cstates);
        }

        if let Some(users) = logged_in_users {
            system_metrics["summary"]["logged_in_users"] = json!(users);
        }

        debug!("System metrics collected: CPU load {:.2}, Memory {:.1}%", 
               cpu_load_5, memory_usage_percent);

        Ok(CollectorOutput {
            agent_type: AgentType::System,
            data: system_metrics,
        })
    }
}