watcheragent/WatcherAgent/src/hardware/disk.rs

use crate::models::DiskInfoDetailed;

use std::error::Error;
use anyhow::Result;
use sysinfo::DiskUsage;
use sysinfo::{Component, Components, Disk, Disks};
use serde::Serialize;


/// # Disk Hardware Module
///
/// This module provides disk information collection for WatcherAgent, including total and per-disk statistics and temperature data.
///
/// ## Responsibilities
/// - **Disk Enumeration:** Lists all physical disks and their properties.
/// - **Usage Calculation:** Computes total and per-disk usage, available space, and usage percentage.
/// - **Temperature Monitoring:** Associates disk components with temperature sensors if available.
///
/// ## Units
/// - All sizes are in **bytes** unless otherwise noted.
/// - Temperatures are in **degrees Celsius (°C)**.
///
/// Summary of disk statistics for the system.
///
/// # Fields
/// - `total_size`: Total disk size in bytes (all disks > 100MB)
/// - `total_used`: Total used disk space in bytes
/// - `total_available`: Total available disk space in bytes
/// - `total_usage`: Usage percentage (0.0–100.0)
/// - `detailed_info`: Vector of [`DiskInfoDetailed`] for each disk
#[derive(Serialize, Debug)]
pub struct DiskInfo {
    pub total_size: Option<f64>,
    pub total_used: Option<f64>,
    pub total_available: Option<f64>,
    pub total_usage: Option<f64>,
    pub detailed_info: Vec<DiskInfoDetailed>,
}

/// Collects disk information for all detected disks, including usage and temperature.
///
/// This function enumerates all disks, calculates usage statistics, and attempts to associate temperature sensors with disk components.
///
/// # Returns
/// * `Result<DiskInfo, Box<dyn std::error::Error + Send + Sync>>` - Disk statistics and details, or error if collection fails.
pub async fn get_disk_info() -> Result<DiskInfo, Box<dyn std::error::Error + Send + Sync>> {
    let disks = Disks::new_with_refreshed_list();
    let mut detailed_info = Vec::new();
    
    // Collect detailed disk information
    for disk in disks.list() {
        if disk.kind() == sysinfo::DiskKind::Unknown(0) {
            continue;
        }
        
        let disk_used = disk.total_space() - disk.available_space();
        detailed_info.push(DiskInfoDetailed {
            disk_name: disk.name().to_string_lossy().into_owned(),
            disk_kind: format!("{:?}", disk.kind()),
            disk_total_space: disk.total_space() as f64,
            disk_available_space: disk.available_space() as f64,
            disk_used_space: disk_used as f64,
            disk_mount_point: disk.mount_point().to_string_lossy().into_owned(),
            component_disk_label: String::new(),
            component_disk_temperature: 0.0,
        });
    }
    
    // Get component temperatures
    let components = Components::new_with_refreshed_list();
    for component in &components {
        if let Some(temperature) = component.temperature() {
            // Update detailed info with temperature data if it matches a disk component
            for disk_info in &mut detailed_info {
                if component.label().contains(&disk_info.disk_name) {
                    disk_info.component_disk_label = component.label().to_string();
                    disk_info.component_disk_temperature = temperature;
                }
            }
        }
    }
    
    // Calculate totals (only disks > 100MB)
    let (total_size, total_used, total_available) = calculate_disk_totals(&disks);
    
    let (total_size, total_used, total_available, total_usage) = if total_size > 0.0 {
        (total_size, total_used, total_available, (total_used / total_size) * 100.0)
    } else {
        match get_disk_info_fallback() {
            Ok(fallback_data) => fallback_data,
            Err(_) => (0.0, 0.0, 0.0, 0.0), // Default values if fallback fails
        }
    };
    
    Ok(DiskInfo {
        total_size: if total_size > 0.0 { Some(total_size) } else { None },
        total_used: if total_used > 0.0 { Some(total_used) } else { None },
        total_available: if total_available > 0.0 { Some(total_available) } else { None },
        total_usage: if total_usage > 0.0 { Some(total_usage) } else { None },
        detailed_info,
    })
}

fn calculate_disk_totals(disks: &Disks) -> (f64, f64, f64) {
    let mut total_size = 0u64;
    let mut total_used = 0u64;
    let mut total_available = 0u64;
    
    for disk in disks.list() {
        if disk.total_space() > 100 * 1024 * 1024 { // > 100MB
            total_size += disk.total_space();
            total_available += disk.available_space();
            total_used += disk.total_space() - disk.available_space();
        }
    }
    
    (total_size as f64, total_used as f64, total_available as f64)
}

#[cfg(target_os = "linux")]
fn get_disk_info_fallback() -> Result<(f64, f64, f64, f64), Box<dyn Error + Send + Sync>> {
    use std::process::Command;
    
    let output = Command::new("df")
        .arg("-B1")
        .arg("--output=size,used,avail")
        .output()?;
        
    let stdout = String::from_utf8_lossy(&output.stdout);
    let mut total_size = 0u64;
    let mut total_used = 0u64;
    let mut total_available = 0u64;
    let mut count = 0;
    
    for line in stdout.lines().skip(1) {
        let parts: Vec<&str> = line.split_whitespace().collect();
        if parts.len() >= 3 {
            if let (Ok(size), Ok(used), Ok(avail)) = (
                parts[0].parse::<u64>(),
                parts[1].parse::<u64>(),
                parts[2].parse::<u64>(),
            ) {
                total_size += size;
                total_used += used;
                total_available += avail;
                count += 1;
            }
        }
    }
    
    let usage = if total_size > 0 {
        (total_used as f64 / total_size as f64) * 100.0
    } else {
        0.0
    };
    
    Ok((total_size as f64, total_used as f64, total_available as f64, usage))
}

#[cfg(not(target_os = "linux"))]
fn get_disk_info_fallback() -> Result<(f64, f64, f64, f64), Box<dyn Error + Send + Sync>> {
    Ok((0.0, 0.0, 0.0, 0.0))
}

pub fn _get_disk_temp_for_component(component: &Component) -> Option<f64> {
    component.temperature().map(|temp| temp as f64)
}

pub fn _get_disk_load_for_disk(disk: &Disk) -> Result<(f64, f64, f64, f64), Box<dyn Error>> {
    let usage: DiskUsage = disk.usage();

    // Assuming DiskUsage has these methods:
    let total_written_bytes = usage.total_written_bytes as f64;
    let written_bytes = usage.written_bytes as f64;
    let total_read_bytes = usage.total_read_bytes as f64;
    let read_bytes = usage.read_bytes as f64;

    Ok((
        total_written_bytes,
        written_bytes,
        total_read_bytes,
        read_bytes,
    ))
}