zinit 0.3.6

//! Disk detection, enumeration, and priority sorting

use super::error::StorageError;
use std::fs;
use std::path::{Path, PathBuf};

/// Type of storage device
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum DiskType {
    /// NVMe device (highest priority)
    Nvme,
    /// SSD (SATA/SAS, non-rotational)
    Ssd,
    /// HDD (rotational disk)
    Hdd,
}

/// Information about a block device
#[derive(Debug, Clone)]
pub struct BlockDevice {
    /// Device name (e.g., "nvme0n1", "sda")
    pub name: String,
    /// Full device path (e.g., "/dev/nvme0n1")
    pub path: PathBuf,
    /// Type of device
    pub disk_type: DiskType,
    /// Size in bytes
    pub size: u64,
}

impl BlockDevice {
    /// Get the sysfs path for this device
    pub fn sysfs_path(&self) -> PathBuf {
        PathBuf::from("/sys/block").join(&self.name)
    }
}

/// Enumerate all candidate block devices, sorted by priority
///
/// Returns devices in priority order: NVMe > SSD > HDD
/// Within each tier, devices are sorted alphabetically.
///
/// Filters out:
/// - loop*, ram*, dm-*, sr*, fd*, zram* devices
/// - Devices with size 0 or unreadable size
pub fn enumerate_disks() -> Result<Vec<BlockDevice>, StorageError> {
    let mut devices = Vec::new();

    let entries = fs::read_dir("/sys/block").map_err(|e| {
        StorageError::EnumerationFailed(format!("failed to read /sys/block: {}", e))
    })?;

    for entry in entries.flatten() {
        let name = entry.file_name().to_string_lossy().to_string();

        // Skip filtered device types
        if should_filter_device(&name) {
            log::debug!("skipping filtered device: {}", name);
            continue;
        }

        // Get device size
        let size = match read_device_size(&name) {
            Some(s) if s > 0 => s,
            _ => {
                log::debug!("skipping device with invalid size: {}", name);
                continue;
            }
        };

        // Determine device type
        let disk_type = determine_disk_type(&name);

        devices.push(BlockDevice {
            path: PathBuf::from("/dev").join(&name),
            name,
            disk_type,
            size,
        });
    }

    // Sort by priority (NVMe > SSD > HDD) then alphabetically
    devices.sort_by(|a, b| {
        a.disk_type
            .cmp(&b.disk_type)
            .then_with(|| a.name.cmp(&b.name))
    });

    log::debug!("found {} candidate disks", devices.len());
    for dev in &devices {
        log::debug!("  {} ({:?}, {} bytes)", dev.name, dev.disk_type, dev.size);
    }

    Ok(devices)
}

/// Check if a device name should be filtered out
fn should_filter_device(name: &str) -> bool {
    // Filter out non-disk devices
    name.starts_with("loop")
        || name.starts_with("ram")
        || name.starts_with("dm-")
        || name.starts_with("sr")
        || name.starts_with("fd")
        || name.starts_with("zram")
}

/// Read device size from sysfs (in bytes)
fn read_device_size(name: &str) -> Option<u64> {
    let size_path = Path::new("/sys/block").join(name).join("size");
    let contents = fs::read_to_string(size_path).ok()?;
    let sectors: u64 = contents.trim().parse().ok()?;
    // Size is in 512-byte sectors
    Some(sectors * 512)
}

/// Determine the type of disk (NVMe, SSD, or HDD)
fn determine_disk_type(name: &str) -> DiskType {
    if name.starts_with("nvme") {
        return DiskType::Nvme;
    }

    // For SATA/SAS devices, check rotational flag
    let rotational_path = Path::new("/sys/block").join(name).join("queue/rotational");

    if let Ok(contents) = fs::read_to_string(rotational_path)
        && contents.trim() == "0"
    {
        return DiskType::Ssd;
    }

    DiskType::Hdd
}

/// Check if a device has any existing partitions
pub fn has_partitions(device: &BlockDevice) -> bool {
    let sysfs = device.sysfs_path();

    // Check for partition directories in sysfs
    if let Ok(entries) = fs::read_dir(&sysfs) {
        for entry in entries.flatten() {
            let name = entry.file_name().to_string_lossy().to_string();
            // Partition directories start with the device name
            if name.starts_with(&device.name) && name != device.name {
                log::debug!("device {} has partition: {}", device.name, name);
                return true;
            }
        }
    }

    false
}

/// Check if a device is currently mounted
pub fn is_mounted(device: &BlockDevice) -> bool {
    if let Ok(mounts) = fs::read_to_string("/proc/mounts") {
        let dev_path = device.path.to_string_lossy();
        for line in mounts.lines() {
            if line.starts_with(&*dev_path) {
                return true;
            }
        }
    }
    false
}

/// Check if a device has holders (in use by device mapper, etc.)
pub fn has_holders(device: &BlockDevice) -> bool {
    let holders_path = device.sysfs_path().join("holders");
    if let Ok(entries) = fs::read_dir(holders_path) {
        let count = entries.count();
        if count > 0 {
            log::debug!("device {} has {} holders", device.name, count);
            return true;
        }
    }
    false
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_filter_devices() {
        assert!(should_filter_device("loop0"));
        assert!(should_filter_device("ram0"));
        assert!(should_filter_device("dm-0"));
        assert!(should_filter_device("sr0"));
        assert!(should_filter_device("fd0"));
        assert!(should_filter_device("zram0"));

        assert!(!should_filter_device("sda"));
        assert!(!should_filter_device("nvme0n1"));
    }

    #[test]
    fn test_disk_type_priority() {
        assert!(DiskType::Nvme < DiskType::Ssd);
        assert!(DiskType::Ssd < DiskType::Hdd);
    }
}