lcpfs 2026.1.102

// Copyright 2025 LunaOS Contributors
// SPDX-License-Identifier: Apache-2.0
//
// RAID-Z Expansion
// Add disks to RAID-Z without full rebuild.

use alloc::collections::BTreeMap;
use alloc::vec::Vec;

/// RAID-Z expansion state
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ExpansionState {
    /// No expansion in progress
    Idle,
    /// Expansion in progress
    Rebalancing,
    /// Expansion paused
    Paused,
    /// Expansion complete
    Complete,
    /// Expansion failed
    Failed,
}

/// RAID-Z expansion metadata
#[derive(Debug, Clone)]
pub struct ExpansionMetadata {
    /// Current expansion state
    pub state: ExpansionState,
    /// Original stripe width
    pub old_width: usize,
    /// New stripe width (after expansion)
    pub new_width: usize,
    /// Number of blocks rebalanced
    pub blocks_rebalanced: u64,
    /// Total blocks to rebalance
    pub total_blocks: u64,
    /// Expansion start timestamp
    pub start_time: u64,
    /// New disk IDs added to pool
    pub new_disks: Vec<u64>,
}

impl ExpansionMetadata {
    /// Create new expansion metadata
    pub fn new(old_width: usize, new_width: usize, new_disks: Vec<u64>) -> Self {
        Self {
            state: ExpansionState::Idle,
            old_width,
            new_width,
            blocks_rebalanced: 0,
            total_blocks: 0,
            start_time: 0,
            new_disks,
        }
    }

    /// Get expansion progress percentage
    pub fn progress(&self) -> f64 {
        if self.total_blocks == 0 {
            return 0.0;
        }
        (self.blocks_rebalanced as f64 / self.total_blocks as f64) * 100.0
    }

    /// Check if expansion is active
    pub fn is_active(&self) -> bool {
        matches!(self.state, ExpansionState::Rebalancing)
    }
}

/// RAID-Z expansion engine
pub struct RaidzExpansion {
    /// Expansion metadata
    metadata: ExpansionMetadata,
    /// Rebalancing priority (blocks per second)
    throttle_bps: u64,
    /// Stripe assignment map (stripe_id -> disk_indices)
    stripe_map: BTreeMap<u64, Vec<usize>>,
}

/// RAID-Z configuration errors
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RaidzConfigError {
    /// Stripe width is too small (minimum 2 for RAID-Z1)
    StripeWidthTooSmall,
    /// Stripe width exceeds maximum (255 disks)
    StripeWidthTooLarge,
    /// No disks provided for expansion
    NoDisksProvided,
    /// Cannot expand by more than 50% at once
    ExpansionTooLarge,
    /// Parity count exceeds data disks
    InvalidParityCount,
}

impl RaidzExpansion {
    /// Default rebalancing throttle (10 MB/s)
    const DEFAULT_THROTTLE: u64 = 10 * 1024 * 1024;

    /// Minimum stripe width (1 data + 1 parity = RAID-Z1)
    const MIN_STRIPE_WIDTH: usize = 2;

    /// Maximum stripe width (practical limit for GF(2^8))
    const MAX_STRIPE_WIDTH: usize = 255;

    /// Maximum expansion ratio (prevent 2-disk → 100-disk expansions)
    const MAX_EXPANSION_RATIO: f64 = 1.5; // 50% growth per expansion

    /// Create new RAID-Z expansion with validation
    ///
    /// # Arguments
    /// * `old_width` - Current stripe width (data + parity disks)
    /// * `new_disks` - New disk IDs to add
    ///
    /// # Returns
    /// * `Ok(RaidzExpansion)` if configuration is valid
    /// * `Err(RaidzConfigError)` if configuration is invalid
    ///
    /// # Example
    /// ```rust,ignore
    /// // Expand 4-disk RAID-Z1 to 5 disks
    /// let exp = RaidzExpansion::try_new(4, vec![100])?;
    /// ```
    pub fn try_new(old_width: usize, new_disks: Vec<u64>) -> Result<Self, RaidzConfigError> {
        // Validate old stripe width
        if old_width < Self::MIN_STRIPE_WIDTH {
            return Err(RaidzConfigError::StripeWidthTooSmall);
        }

        // Validate new disks provided
        if new_disks.is_empty() {
            return Err(RaidzConfigError::NoDisksProvided);
        }

        let new_width = old_width + new_disks.len();

        // Validate new stripe width
        if new_width > Self::MAX_STRIPE_WIDTH {
            return Err(RaidzConfigError::StripeWidthTooLarge);
        }

        // Validate expansion ratio (prevent runaway expansions)
        let expansion_ratio = new_width as f64 / old_width as f64;
        if expansion_ratio > Self::MAX_EXPANSION_RATIO {
            return Err(RaidzConfigError::ExpansionTooLarge);
        }

        Ok(Self {
            metadata: ExpansionMetadata::new(old_width, new_width, new_disks),
            throttle_bps: Self::DEFAULT_THROTTLE,
            stripe_map: BTreeMap::new(),
        })
    }

    /// Create new RAID-Z expansion (legacy, panics on invalid config)
    ///
    /// # Arguments
    /// * `old_width` - Current stripe width (data + parity disks)
    /// * `new_disks` - New disk IDs to add
    ///
    /// # Panics
    /// Panics if configuration is invalid. Prefer `try_new()` for production code.
    pub fn new(old_width: usize, new_disks: Vec<u64>) -> Self {
        Self::try_new(old_width, new_disks).expect("Invalid RAID-Z expansion configuration")
    }

    /// Start expansion process
    ///
    /// # Arguments
    /// * `total_blocks` - Total number of blocks to rebalance
    /// * `start_time` - Start timestamp
    pub fn start(&mut self, total_blocks: u64, start_time: u64) {
        self.metadata.state = ExpansionState::Rebalancing;
        self.metadata.total_blocks = total_blocks;
        self.metadata.start_time = start_time;
        self.metadata.blocks_rebalanced = 0;
    }

    /// Pause expansion
    pub fn pause(&mut self) {
        if self.metadata.is_active() {
            self.metadata.state = ExpansionState::Paused;
        }
    }

    /// Resume expansion
    pub fn resume(&mut self) {
        if self.metadata.state == ExpansionState::Paused {
            self.metadata.state = ExpansionState::Rebalancing;
        }
    }

    /// Rebalance a batch of stripes
    ///
    /// # Arguments
    /// * `stripe_ids` - Stripe IDs to rebalance
    ///
    /// # Returns
    /// * `Ok(blocks_moved)` - Number of blocks successfully rebalanced
    /// * `Err(msg)` - Error message
    ///
    /// # Algorithm
    /// 1. For each stripe:
    ///    a. Read all data blocks
    ///    b. Recalculate parity with new stripe width
    ///    c. Write data to new stripe layout
    ///    d. Verify checksums
    /// 2. Update stripe map with new disk assignments
    pub fn rebalance_stripes(&mut self, stripe_ids: &[u64]) -> Result<u64, &'static str> {
        if !self.metadata.is_active() {
            return Err("Expansion not active");
        }

        let mut blocks_moved = 0;

        for &stripe_id in stripe_ids {
            // Simulate rebalancing (real implementation would read/write blocks)
            // New stripe layout spreads data across more disks
            let new_layout = self.calculate_new_stripe_layout(stripe_id);

            // Update stripe map
            self.stripe_map.insert(stripe_id, new_layout);

            // Assume each stripe has (old_width - parity) data blocks
            blocks_moved += (self.metadata.old_width - 1) as u64;
        }

        self.metadata.blocks_rebalanced += blocks_moved;

        // Check if expansion is complete
        if self.metadata.blocks_rebalanced >= self.metadata.total_blocks {
            self.metadata.state = ExpansionState::Complete;
        }

        Ok(blocks_moved)
    }

    /// Calculate new stripe layout for a stripe
    ///
    /// Distributes data blocks across old and new disks evenly
    ///
    /// # Arguments
    /// * `stripe_id` - Stripe identifier
    ///
    /// # Returns
    /// Vector of disk indices for this stripe
    fn calculate_new_stripe_layout(&self, stripe_id: u64) -> Vec<usize> {
        let mut layout = Vec::new();

        // Round-robin distribution across all disks
        for i in 0..self.metadata.new_width {
            let disk_idx = ((stripe_id as usize) + i) % self.metadata.new_width;
            layout.push(disk_idx);
        }

        layout
    }

    /// Set rebalancing throttle
    ///
    /// # Arguments
    /// * `bytes_per_second` - Maximum bytes to rebalance per second
    pub fn set_throttle(&mut self, bytes_per_second: u64) {
        self.throttle_bps = bytes_per_second;
    }

    /// Get expansion metadata
    pub fn metadata(&self) -> &ExpansionMetadata {
        &self.metadata
    }

    /// Estimate time remaining
    ///
    /// # Arguments
    /// * `current_time` - Current timestamp
    ///
    /// # Returns
    /// Estimated seconds until completion
    pub fn estimate_time_remaining(&self, current_time: u64) -> u64 {
        if !self.metadata.is_active() || self.metadata.blocks_rebalanced == 0 {
            return 0;
        }

        let elapsed = current_time - self.metadata.start_time;
        let blocks_per_second = self.metadata.blocks_rebalanced / elapsed;

        let remaining_blocks = self.metadata.total_blocks - self.metadata.blocks_rebalanced;
        if blocks_per_second > 0 {
            remaining_blocks / blocks_per_second
        } else {
            u64::MAX // Unknown
        }
    }
}

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

    #[test]
    fn test_expansion_metadata() {
        let meta = ExpansionMetadata::new(5, 7, vec![100, 101]);

        assert_eq!(meta.old_width, 5);
        assert_eq!(meta.new_width, 7);
        assert_eq!(meta.state, ExpansionState::Idle);
        assert_eq!(meta.progress(), 0.0);
        assert!(!meta.is_active());
    }

    #[test]
    fn test_expansion_progress() {
        let mut meta = ExpansionMetadata::new(4, 6, vec![200]);
        meta.total_blocks = 1000;
        meta.blocks_rebalanced = 250;

        assert_eq!(meta.progress(), 25.0);
    }

    #[test]
    fn test_expansion_start() {
        let mut exp = RaidzExpansion::new(4, vec![100]);
        exp.start(1000, 12345);

        assert!(exp.metadata.is_active());
        assert_eq!(exp.metadata.state, ExpansionState::Rebalancing);
        assert_eq!(exp.metadata.total_blocks, 1000);
        assert_eq!(exp.metadata.start_time, 12345);
    }

    #[test]
    fn test_expansion_pause_resume() {
        let mut exp = RaidzExpansion::new(4, vec![100]);
        exp.start(1000, 0);

        exp.pause();
        assert_eq!(exp.metadata.state, ExpansionState::Paused);

        exp.resume();
        assert!(exp.metadata.is_active());
    }

    #[test]
    fn test_rebalance_stripes() {
        let mut exp = RaidzExpansion::new(4, vec![100, 101]);
        exp.start(1000, 0);

        let stripes = vec![1, 2, 3];
        let blocks_moved = exp
            .rebalance_stripes(&stripes)
            .expect("test: operation should succeed");

        // Each stripe has 3 data blocks (4 - 1 parity)
        assert_eq!(blocks_moved, 9); // 3 stripes × 3 blocks
        assert_eq!(exp.metadata.blocks_rebalanced, 9);
    }

    #[test]
    fn test_expansion_complete() {
        let mut exp = RaidzExpansion::new(3, vec![100]);
        exp.start(10, 0);

        // Rebalance enough to complete
        let stripes = vec![1, 2, 3, 4, 5];
        exp.rebalance_stripes(&stripes)
            .expect("test: operation should succeed");

        assert_eq!(exp.metadata.state, ExpansionState::Complete);
    }

    #[test]
    fn test_new_stripe_layout() {
        let exp = RaidzExpansion::new(4, vec![100, 101]);

        let layout1 = exp.calculate_new_stripe_layout(0);
        let layout2 = exp.calculate_new_stripe_layout(1);

        // Should distribute across all 6 disks
        assert_eq!(layout1.len(), 6);
        assert_eq!(layout2.len(), 6);

        // Layouts should be different (rotated)
        assert_ne!(layout1, layout2);
    }

    #[test]
    fn test_time_estimate() {
        let mut exp = RaidzExpansion::new(4, vec![100]);
        exp.start(1000, 0);

        // Rebalance 3 stripes × 3 blocks/stripe = 9 blocks
        exp.rebalance_stripes(&[1, 2, 3])
            .expect("test: operation should succeed");

        // After 9 seconds, estimate time remaining
        let remaining = exp.estimate_time_remaining(9);

        // 9 blocks in 9s = 1 block/s
        // 991 blocks remaining / 1 block/s = 991s
        assert_eq!(remaining, 991);
    }

    #[test]
    fn test_throttle() {
        let mut exp = RaidzExpansion::new(4, vec![100]);
        exp.set_throttle(5 * 1024 * 1024); // 5 MB/s

        assert_eq!(exp.throttle_bps, 5 * 1024 * 1024);
    }

    #[test]
    fn test_config_validation_stripe_too_small() {
        let result = RaidzExpansion::try_new(1, vec![100]);
        assert_eq!(result.err(), Some(RaidzConfigError::StripeWidthTooSmall));
    }

    #[test]
    fn test_config_validation_no_disks() {
        let result = RaidzExpansion::try_new(4, vec![]);
        assert_eq!(result.err(), Some(RaidzConfigError::NoDisksProvided));
    }

    #[test]
    fn test_config_validation_expansion_too_large() {
        // Try to expand from 4 disks to 8 (100% growth, exceeds 50% limit)
        let result = RaidzExpansion::try_new(4, vec![100, 101, 102, 103]);
        assert_eq!(result.err(), Some(RaidzConfigError::ExpansionTooLarge));
    }

    #[test]
    fn test_config_validation_valid() {
        // Valid expansion: 4 disks → 5 disks (25% growth)
        let result = RaidzExpansion::try_new(4, vec![100]);
        assert!(result.is_ok());

        // Valid expansion: 4 disks → 6 disks (50% growth, at limit)
        let result = RaidzExpansion::try_new(4, vec![100, 101]);
        assert!(result.is_ok());
    }
}