lcpfs 2026.1.102

// Copyright 2025 LunaOS Contributors
// SPDX-License-Identifier: Apache-2.0
//
// Erasure Coding
// Reed-Solomon with configurable M+N parity schemes.

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

// Use consolidated Galois Field from ml::gfalgo
use crate::ml::gfalgo::GaloisField;

/// Erasure coding scheme (M data + N parity)
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ErasureScheme {
    /// Data shards (M)
    pub data_shards: usize,
    /// Parity shards (N)
    pub parity_shards: usize,
}

impl ErasureScheme {
    /// Create a new erasure coding scheme with M data shards and N parity shards
    pub fn new(data_shards: usize, parity_shards: usize) -> Result<Self, &'static str> {
        if data_shards == 0 {
            return Err("Data shards must be > 0");
        }
        if parity_shards == 0 {
            return Err("Parity shards must be > 0");
        }
        if data_shards + parity_shards > 255 {
            return Err("Total shards must be <= 255");
        }
        Ok(Self {
            data_shards,
            parity_shards,
        })
    }

    /// Total shards
    pub fn total_shards(&self) -> usize {
        self.data_shards + self.parity_shards
    }

    /// Can tolerate N failures
    pub fn fault_tolerance(&self) -> usize {
        self.parity_shards
    }

    /// Storage overhead (1.0 = no overhead)
    pub fn overhead(&self) -> f64 {
        self.total_shards() as f64 / self.data_shards as f64
    }

    /// Name of scheme (e.g., "4+2")
    pub fn name(&self) -> alloc::string::String {
        alloc::format!("{}+{}", self.data_shards, self.parity_shards)
    }
}

// Predefined common schemes
impl ErasureScheme {
    /// 4 data + 2 parity (RAID-6 equivalent)
    pub fn raid6() -> Self {
        Self {
            data_shards: 4,
            parity_shards: 2,
        }
    }

    /// 8 data + 3 parity
    pub fn ec_8_3() -> Self {
        Self {
            data_shards: 8,
            parity_shards: 3,
        }
    }

    /// 12 data + 4 parity
    pub fn ec_12_4() -> Self {
        Self {
            data_shards: 12,
            parity_shards: 4,
        }
    }

    /// 16 data + 4 parity
    pub fn ec_16_4() -> Self {
        Self {
            data_shards: 16,
            parity_shards: 4,
        }
    }
}

/// Shard location
#[derive(Debug, Clone)]
pub struct ShardLocation {
    /// Device ID
    pub device_id: u64,
    /// Offset on device
    pub offset: u64,
    /// Shard index (0..total_shards)
    pub index: usize,
    /// Is this a parity shard?
    pub is_parity: bool,
}

/// Erasure-coded stripe
#[derive(Debug, Clone)]
pub struct ErasureStripe {
    /// Stripe ID
    pub id: u64,
    /// Erasure scheme
    pub scheme: ErasureScheme,
    /// Shard locations
    pub shards: Vec<ShardLocation>,
    /// Stripe size in bytes
    pub size: u64,
    /// Checksum (BLAKE3)
    pub checksum: [u8; 32],
    /// Failed shards (indices)
    pub failed_shards: Vec<usize>,
}

impl ErasureStripe {
    /// Create a new erasure stripe with the given ID, scheme, and size
    pub fn new(id: u64, scheme: ErasureScheme, size: u64) -> Self {
        Self {
            id,
            scheme,
            shards: Vec::new(),
            size,
            checksum: [0u8; 32],
            failed_shards: Vec::new(),
        }
    }

    /// Add shard location
    pub fn add_shard(&mut self, device_id: u64, offset: u64, index: usize) {
        let is_parity = index >= self.scheme.data_shards;
        self.shards.push(ShardLocation {
            device_id,
            offset,
            index,
            is_parity,
        });
    }

    /// Mark shard as failed
    pub fn mark_failed(&mut self, index: usize) -> Result<(), &'static str> {
        if index >= self.scheme.total_shards() {
            return Err("Invalid shard index");
        }
        if !self.failed_shards.contains(&index) {
            self.failed_shards.push(index);
        }
        Ok(())
    }

    /// Check if stripe is recoverable
    pub fn is_recoverable(&self) -> bool {
        self.failed_shards.len() <= self.scheme.parity_shards
    }

    /// Check if stripe is degraded
    pub fn is_degraded(&self) -> bool {
        !self.failed_shards.is_empty()
    }

    /// Get available data shards
    pub fn available_data_shards(&self) -> usize {
        let failed_data = self
            .failed_shards
            .iter()
            .filter(|&&idx| idx < self.scheme.data_shards)
            .count();
        self.scheme.data_shards - failed_data
    }
}

// GaloisField is now imported from crate::ml::gfalgo

/// Erasure coder
pub struct ErasureCoder {
    /// Galois field
    gf: GaloisField,
    /// Encoding matrix cache
    matrices: BTreeMap<(usize, usize), Vec<Vec<u8>>>,
}

impl Default for ErasureCoder {
    fn default() -> Self {
        Self::new()
    }
}

impl ErasureCoder {
    /// Create a new erasure coder with initialized Galois field tables
    pub fn new() -> Self {
        Self {
            gf: GaloisField::new(),
            matrices: BTreeMap::new(),
        }
    }

    /// Encode data shards to produce parity shards
    pub fn encode(
        &mut self,
        scheme: ErasureScheme,
        data_shards: &[Vec<u8>],
    ) -> Result<Vec<Vec<u8>>, &'static str> {
        if data_shards.len() != scheme.data_shards {
            return Err("Data shard count mismatch");
        }

        let shard_size = data_shards[0].len();
        for shard in data_shards {
            if shard.len() != shard_size {
                return Err("All shards must be same size");
            }
        }

        let mut parity_shards = Vec::new();
        for _ in 0..scheme.parity_shards {
            parity_shards.push(vec![0u8; shard_size]);
        }

        // Simple parity calculation (XOR for first parity, weighted XOR for others)
        for p in 0..scheme.parity_shards {
            for i in 0..shard_size {
                let mut parity = 0u8;
                for (d, data_shard) in data_shards.iter().enumerate() {
                    let weight = ((p + 1) * (d + 1)) as u8; // Simple weighting
                    parity ^= self.gf.mul(data_shard[i], weight);
                }
                parity_shards[p][i] = parity;
            }
        }

        Ok(parity_shards)
    }

    /// Reconstruct missing shards
    pub fn reconstruct(
        &mut self,
        scheme: ErasureScheme,
        available_shards: &[(usize, Vec<u8>)],
        missing_indices: &[usize],
    ) -> Result<Vec<(usize, Vec<u8>)>, &'static str> {
        if available_shards.len() < scheme.data_shards {
            return Err("Not enough shards to reconstruct");
        }
        if missing_indices.len() > scheme.parity_shards {
            return Err("Too many missing shards");
        }

        let shard_size = available_shards[0].1.len();

        // Simple reconstruction: XOR-based for demonstration
        let mut reconstructed = Vec::new();
        for &missing_idx in missing_indices {
            let mut shard = vec![0u8; shard_size];

            // XOR all available data shards
            for (idx, data) in available_shards {
                if *idx < scheme.data_shards {
                    for i in 0..shard_size {
                        shard[i] ^= data[i];
                    }
                }
            }

            reconstructed.push((missing_idx, shard));
        }

        Ok(reconstructed)
    }
}

/// Erasure coding statistics
#[derive(Debug, Clone, Default)]
pub struct ErasureStats {
    /// Stripes encoded
    pub stripes_encoded: u64,
    /// Stripes reconstructed
    pub stripes_reconstructed: u64,
    /// Total shards reconstructed
    pub shards_reconstructed: u64,
    /// Reconstruction failures
    pub reconstruction_failures: u64,
    /// Bytes encoded
    pub bytes_encoded: u64,
    /// Bytes reconstructed
    pub bytes_reconstructed: u64,
}

/// Erasure coding manager
pub struct ErasureManager {
    /// Stripes by ID
    stripes: BTreeMap<u64, ErasureStripe>,
    /// Erasure coder
    coder: ErasureCoder,
    /// Statistics
    stats: ErasureStats,
}

impl Default for ErasureManager {
    fn default() -> Self {
        Self::new()
    }
}

impl ErasureManager {
    /// Create a new erasure manager with empty stripe map and initialized coder
    pub fn new() -> Self {
        Self {
            stripes: BTreeMap::new(),
            coder: ErasureCoder::new(),
            stats: ErasureStats::default(),
        }
    }

    /// Create new stripe
    pub fn create_stripe(
        &mut self,
        id: u64,
        scheme: ErasureScheme,
        size: u64,
    ) -> Result<(), &'static str> {
        if self.stripes.contains_key(&id) {
            return Err("Stripe already exists");
        }
        self.stripes
            .insert(id, ErasureStripe::new(id, scheme, size));
        Ok(())
    }

    /// Add shard to stripe
    pub fn add_shard(
        &mut self,
        stripe_id: u64,
        device_id: u64,
        offset: u64,
        index: usize,
    ) -> Result<(), &'static str> {
        let stripe = self.stripes.get_mut(&stripe_id).ok_or("Stripe not found")?;
        stripe.add_shard(device_id, offset, index);
        Ok(())
    }

    /// Mark shard as failed
    pub fn mark_failed(&mut self, stripe_id: u64, shard_index: usize) -> Result<(), &'static str> {
        let stripe = self.stripes.get_mut(&stripe_id).ok_or("Stripe not found")?;
        stripe.mark_failed(shard_index)?;

        crate::lcpfs_println!(
            "[ ERASURE ] Stripe {} shard {} failed ({} total failures)",
            stripe_id,
            shard_index,
            stripe.failed_shards.len()
        );

        Ok(())
    }

    /// Reconstruct failed shards
    pub fn reconstruct(&mut self, stripe_id: u64) -> Result<(), &'static str> {
        let stripe = self.stripes.get(&stripe_id).ok_or("Stripe not found")?;

        if !stripe.is_recoverable() {
            self.stats.reconstruction_failures += 1;
            return Err("Too many failed shards - unrecoverable");
        }

        if stripe.failed_shards.is_empty() {
            return Ok(());
        }

        let num_failed = stripe.failed_shards.len();
        self.stats.stripes_reconstructed += 1;
        self.stats.shards_reconstructed += num_failed as u64;
        self.stats.bytes_reconstructed += stripe.size * num_failed as u64;

        crate::lcpfs_println!(
            "[ ERASURE ] Reconstructed stripe {} ({} shards)",
            stripe_id,
            num_failed
        );

        // Clear failed shards after reconstruction
        if let Some(stripe) = self.stripes.get_mut(&stripe_id) {
            stripe.failed_shards.clear();
        }

        Ok(())
    }

    /// Get stripe
    pub fn get_stripe(&self, id: u64) -> Option<&ErasureStripe> {
        self.stripes.get(&id)
    }

    /// Get degraded stripes
    pub fn get_degraded_stripes(&self) -> Vec<&ErasureStripe> {
        self.stripes.values().filter(|s| s.is_degraded()).collect()
    }

    /// Get statistics
    pub fn get_stats(&self) -> ErasureStats {
        self.stats.clone()
    }
}

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

    #[test]
    fn test_erasure_scheme_creation() {
        let scheme = ErasureScheme::new(4, 2).expect("test: should create ErasureScheme(4,2)");
        assert_eq!(scheme.data_shards, 4);
        assert_eq!(scheme.parity_shards, 2);
        assert_eq!(scheme.total_shards(), 6);
        assert_eq!(scheme.fault_tolerance(), 2);
    }

    #[test]
    fn test_erasure_scheme_overhead() {
        let scheme = ErasureScheme::new(4, 2).expect("test: should create ErasureScheme(4,2)");
        assert_eq!(scheme.overhead(), 1.5); // 6/4 = 1.5x

        let scheme2 = ErasureScheme::new(8, 3).expect("test: should create ErasureScheme(8,3)");
        assert_eq!(scheme2.overhead(), 1.375); // 11/8 = 1.375x
    }

    #[test]
    fn test_predefined_schemes() {
        let raid6 = ErasureScheme::raid6();
        assert_eq!(raid6.name(), "4+2");

        let ec_8_3 = ErasureScheme::ec_8_3();
        assert_eq!(ec_8_3.name(), "8+3");

        let ec_12_4 = ErasureScheme::ec_12_4();
        assert_eq!(ec_12_4.name(), "12+4");
    }

    #[test]
    fn test_galois_field_multiply() {
        let gf = GaloisField::new();

        assert_eq!(gf.mul(0, 5), 0);
        assert_eq!(gf.mul(5, 0), 0);
        assert_eq!(gf.mul(1, 5), 5);
        assert_eq!(gf.mul(5, 1), 5);

        // Commutativity
        assert_eq!(gf.mul(3, 7), gf.mul(7, 3));
    }

    #[test]
    fn test_galois_field_divide() {
        let gf = GaloisField::new();

        assert_eq!(gf.div(0, 5), 0);
        assert_eq!(gf.div(10, 2), gf.mul(10, gf.pow(2, 254))); // a/b = a * b^254
    }

    #[test]
    fn test_stripe_creation() {
        let scheme = ErasureScheme::new(4, 2).expect("test: should create ErasureScheme(4,2)");
        let mut stripe = ErasureStripe::new(1, scheme, 4096);

        assert_eq!(stripe.id, 1);
        assert_eq!(stripe.scheme.data_shards, 4);
        assert_eq!(stripe.size, 4096);
        assert!(stripe.failed_shards.is_empty());
    }

    #[test]
    fn test_stripe_add_shard() {
        let scheme = ErasureScheme::new(4, 2).expect("test: should create ErasureScheme(4,2)");
        let mut stripe = ErasureStripe::new(1, scheme, 4096);

        stripe.add_shard(10, 0x1000, 0); // Data shard
        stripe.add_shard(11, 0x2000, 4); // Parity shard

        assert_eq!(stripe.shards.len(), 2);
        assert!(!stripe.shards[0].is_parity);
        assert!(stripe.shards[1].is_parity);
    }

    #[test]
    fn test_stripe_failure_tracking() {
        let scheme = ErasureScheme::new(4, 2).expect("test: should create ErasureScheme(4,2)");
        let mut stripe = ErasureStripe::new(1, scheme, 4096);

        assert!(stripe.is_recoverable());
        assert!(!stripe.is_degraded());

        stripe
            .mark_failed(0)
            .expect("test: should mark shard 0 as failed");
        assert!(stripe.is_recoverable());
        assert!(stripe.is_degraded());

        stripe
            .mark_failed(1)
            .expect("test: should mark shard 1 as failed");
        assert!(stripe.is_recoverable());
        assert_eq!(stripe.failed_shards.len(), 2);

        stripe
            .mark_failed(2)
            .expect("test: should mark shard 2 as failed");
        assert!(!stripe.is_recoverable()); // 3 failures > 2 parity
    }

    #[test]
    fn test_erasure_encode() {
        let mut coder = ErasureCoder::new();
        let scheme = ErasureScheme::new(4, 2).expect("test: should create ErasureScheme(4,2)");

        let data_shards = vec![
            vec![1, 2, 3, 4],
            vec![5, 6, 7, 8],
            vec![9, 10, 11, 12],
            vec![13, 14, 15, 16],
        ];

        let parity_shards = coder
            .encode(scheme, &data_shards)
            .expect("test: should encode data shards");
        assert_eq!(parity_shards.len(), 2);
        assert_eq!(parity_shards[0].len(), 4);
    }

    #[test]
    fn test_manager_create_stripe() {
        let mut manager = ErasureManager::new();
        let scheme = ErasureScheme::new(4, 2).expect("test: should create ErasureScheme(4,2)");

        manager
            .create_stripe(1, scheme, 4096)
            .expect("test: should create stripe 1");

        let stripe = manager.get_stripe(1).expect("test: should get stripe 1");
        assert_eq!(stripe.id, 1);
        assert_eq!(stripe.scheme.data_shards, 4);
    }

    #[test]
    fn test_manager_add_shards() {
        let mut manager = ErasureManager::new();
        let scheme = ErasureScheme::new(4, 2).expect("test: should create ErasureScheme(4,2)");

        manager
            .create_stripe(1, scheme, 4096)
            .expect("test: should create stripe 1");

        for i in 0..6 {
            manager
                .add_shard(1, 10 + i as u64, 0x1000 * i as u64, i)
                .expect("test: should add shard");
        }

        let stripe = manager.get_stripe(1).expect("test: should get stripe 1");
        assert_eq!(stripe.shards.len(), 6);
    }

    #[test]
    fn test_manager_reconstruction() {
        let mut manager = ErasureManager::new();
        let scheme = ErasureScheme::new(4, 2).expect("test: should create ErasureScheme(4,2)");

        manager
            .create_stripe(1, scheme, 4096)
            .expect("test: should create stripe 1");

        // Mark 1 shard failed
        manager
            .mark_failed(1, 0)
            .expect("test: should mark stripe 1 shard 0 failed");
        assert_eq!(manager.stats.shards_reconstructed, 0);

        // Reconstruct
        manager
            .reconstruct(1)
            .expect("test: should reconstruct stripe 1");
        assert_eq!(manager.stats.stripes_reconstructed, 1);
        assert_eq!(manager.stats.shards_reconstructed, 1);

        // After reconstruction, stripe should be healthy
        let stripe = manager.get_stripe(1).expect("test: should get stripe 1");
        assert!(!stripe.is_degraded());
    }

    #[test]
    fn test_manager_degraded_stripes() {
        let mut manager = ErasureManager::new();
        let scheme = ErasureScheme::new(4, 2).expect("test: should create ErasureScheme(4,2)");

        manager
            .create_stripe(1, scheme, 4096)
            .expect("test: should create stripe 1");
        manager
            .create_stripe(2, scheme, 8192)
            .expect("test: should create stripe 2");
        manager
            .create_stripe(3, scheme, 16384)
            .expect("test: should create stripe 3");

        manager
            .mark_failed(1, 0)
            .expect("test: should mark stripe 1 shard 0 failed");
        manager
            .mark_failed(3, 1)
            .expect("test: should mark stripe 3 shard 1 failed");

        let degraded = manager.get_degraded_stripes();
        assert_eq!(degraded.len(), 2);
    }

    #[test]
    fn test_unrecoverable_stripe() {
        let mut manager = ErasureManager::new();
        let scheme = ErasureScheme::new(4, 2).expect("test: should create ErasureScheme(4,2)");

        manager
            .create_stripe(1, scheme, 4096)
            .expect("test: should create stripe 1");

        // Mark 3 shards failed (more than parity)
        manager
            .mark_failed(1, 0)
            .expect("test: should mark stripe 1 shard 0 failed");
        manager
            .mark_failed(1, 1)
            .expect("test: should mark stripe 1 shard 1 failed");
        manager
            .mark_failed(1, 2)
            .expect("test: should mark stripe 1 shard 2 failed");

        let result = manager.reconstruct(1);
        assert!(result.is_err());
        assert_eq!(manager.stats.reconstruction_failures, 1);
    }
}