zinit 0.3.7

//! Partitioning via sgdisk

use super::boot::BootMode;
use super::cmd;
use super::detect::BlockDevice;
use super::error::StorageError;
use std::path::PathBuf;

/// Partition information after creation
#[derive(Debug, Clone)]
pub struct PartitionInfo {
    /// Boot partition (ESP for EFI, BIOS boot partition for BIOS)
    pub boot: PartitionDevice,
    /// Data partition (btrfs)
    pub data: PartitionDevice,
}

/// A created partition device
#[derive(Debug, Clone)]
pub struct PartitionDevice {
    /// Partition number
    pub number: u32,
    /// Full device path (e.g., /dev/nvme0n1p1 or /dev/sda1)
    pub path: PathBuf,
    /// Partition name/label
    pub name: String,
}

/// Create the partition table based on boot mode
///
/// For EFI mode:
/// - Partition 1: 512MB ESP (EF00, FAT32)
/// - Partition 2: mossize GB data (8300, btrfs)
///
/// For BIOS mode:
/// - Partition 1: 1MB BIOS boot (EF02, no filesystem)
/// - Partition 2: mossize GB data (8300, btrfs)
pub fn create_partitions(
    device: &BlockDevice,
    boot_mode: BootMode,
    data_size_gb: u32,
) -> Result<PartitionInfo, StorageError> {
    let dev_path = device.path.to_string_lossy();
    log::info!("creating partition table on {}", dev_path);

    // Zap and clear any existing partition table
    zap_disk(device)?;

    match boot_mode {
        BootMode::Efi => create_efi_partitions(device, data_size_gb),
        BootMode::Bios => create_bios_partitions(device, data_size_gb),
    }
}

/// Zap all partition data from disk
fn zap_disk(device: &BlockDevice) -> Result<(), StorageError> {
    let dev_path = device.path.to_string_lossy();

    cmd::run("sgdisk", &["--zap-all", &dev_path]).map_err(|e| {
        StorageError::PartitioningFailed {
            device: device.path.clone(),
            error: format!("zap failed: {}", e),
        }
    })?;

    cmd::run("sgdisk", &["--clear", &dev_path]).map_err(|e| StorageError::PartitioningFailed {
        device: device.path.clone(),
        error: format!("clear failed: {}", e),
    })?;

    Ok(())
}

/// Create EFI partition scheme
fn create_efi_partitions(
    device: &BlockDevice,
    data_size_gb: u32,
) -> Result<PartitionInfo, StorageError> {
    let dev_path = device.path.to_string_lossy();

    // Partition 1: ESP (512MB starting at sector 2048)
    log::info!("creating ESP partition (512MB)");
    sgdisk_create_partition(device, 1, "2048", "+512M")?;
    sgdisk_set_type(device, 1, "EF00")?;
    sgdisk_set_name(device, 1, "mosboot")?;

    // Partition 2: Data
    log::info!("creating data partition ({}GB)", data_size_gb);
    sgdisk_create_partition(device, 2, "0", &format!("+{}G", data_size_gb))?;
    sgdisk_set_type(device, 2, "8300")?;
    sgdisk_set_name(device, 2, "mosdata")?;

    // Wait for kernel to update device nodes
    settle_partitions()?;

    Ok(PartitionInfo {
        boot: PartitionDevice {
            number: 1,
            path: partition_path(&dev_path, 1),
            name: "mosboot".to_string(),
        },
        data: PartitionDevice {
            number: 2,
            path: partition_path(&dev_path, 2),
            name: "mosdata".to_string(),
        },
    })
}

/// Create BIOS partition scheme
fn create_bios_partitions(
    device: &BlockDevice,
    data_size_gb: u32,
) -> Result<PartitionInfo, StorageError> {
    let dev_path = device.path.to_string_lossy();

    // Partition 1: BIOS boot (1MB starting at sector 2048)
    log::info!("creating BIOS boot partition (1MB)");
    sgdisk_create_partition(device, 1, "2048", "+1M")?;
    sgdisk_set_type(device, 1, "EF02")?;
    sgdisk_set_name(device, 1, "mosboot")?;

    // Partition 2: Data
    log::info!("creating data partition ({}GB)", data_size_gb);
    sgdisk_create_partition(device, 2, "0", &format!("+{}G", data_size_gb))?;
    sgdisk_set_type(device, 2, "8300")?;
    sgdisk_set_name(device, 2, "mosdata")?;

    // Wait for kernel to update device nodes
    settle_partitions()?;

    Ok(PartitionInfo {
        boot: PartitionDevice {
            number: 1,
            path: partition_path(&dev_path, 1),
            name: "mosboot".to_string(),
        },
        data: PartitionDevice {
            number: 2,
            path: partition_path(&dev_path, 2),
            name: "mosdata".to_string(),
        },
    })
}

/// Create a partition using sgdisk
fn sgdisk_create_partition(
    device: &BlockDevice,
    number: u32,
    start: &str,
    size: &str,
) -> Result<(), StorageError> {
    let dev_path = device.path.to_string_lossy();
    let new_arg = format!("--new={}:{}:{}", number, start, size);

    cmd::run("sgdisk", &[&new_arg, &dev_path]).map_err(|e| StorageError::PartitioningFailed {
        device: device.path.clone(),
        error: format!("create partition {} failed: {}", number, e),
    })?;

    Ok(())
}

/// Set partition type code
fn sgdisk_set_type(device: &BlockDevice, number: u32, type_code: &str) -> Result<(), StorageError> {
    let dev_path = device.path.to_string_lossy();
    let type_arg = format!("--typecode={}:{}", number, type_code);

    cmd::run("sgdisk", &[&type_arg, &dev_path]).map_err(|e| StorageError::PartitioningFailed {
        device: device.path.clone(),
        error: format!("set type {} failed: {}", number, e),
    })?;

    Ok(())
}

/// Set partition name
fn sgdisk_set_name(device: &BlockDevice, number: u32, name: &str) -> Result<(), StorageError> {
    let dev_path = device.path.to_string_lossy();
    let name_arg = format!("--change-name={}:{}", number, name);

    cmd::run("sgdisk", &[&name_arg, &dev_path]).map_err(|e| StorageError::PartitioningFailed {
        device: device.path.clone(),
        error: format!("set name {} failed: {}", number, e),
    })?;

    Ok(())
}

/// Wait for partition device nodes to appear
fn settle_partitions() -> Result<(), StorageError> {
    // Try udevadm settle, or partprobe, or just sleep
    if cmd::run_allow_fail("udevadm", &["settle", "--timeout=5"]).is_some() {
        return Ok(());
    }

    if cmd::run_allow_fail("partprobe", &[]).is_some() {
        return Ok(());
    }

    // Fallback: just wait a bit
    std::thread::sleep(std::time::Duration::from_secs(1));
    Ok(())
}

/// Get the partition device path
///
/// Handles both nvme (nvme0n1p1) and sd (sda1) naming conventions
fn partition_path(device_path: &str, number: u32) -> PathBuf {
    if device_path.contains("nvme") || device_path.contains("mmcblk") {
        // NVMe and MMC use "p" prefix for partition number
        PathBuf::from(format!("{}p{}", device_path, number))
    } else {
        // SATA/SAS just append the number
        PathBuf::from(format!("{}{}", device_path, number))
    }
}

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

    #[test]
    fn test_partition_path_nvme() {
        assert_eq!(
            partition_path("/dev/nvme0n1", 1),
            PathBuf::from("/dev/nvme0n1p1")
        );
        assert_eq!(
            partition_path("/dev/nvme0n1", 2),
            PathBuf::from("/dev/nvme0n1p2")
        );
    }

    #[test]
    fn test_partition_path_sata() {
        assert_eq!(partition_path("/dev/sda", 1), PathBuf::from("/dev/sda1"));
        assert_eq!(partition_path("/dev/sdb", 2), PathBuf::from("/dev/sdb2"));
    }

    #[test]
    fn test_partition_path_mmcblk() {
        assert_eq!(
            partition_path("/dev/mmcblk0", 1),
            PathBuf::from("/dev/mmcblk0p1")
        );
    }
}