//! Erasure coding module using Reed-Solomon
//!
//! This module provides a wrapper around the `reed-solomon-erasure` crate
//! for encoding and decoding data using Reed-Solomon erasure codes.

use reed_solomon_erasure::galois_8::ReedSolomon;
use thiserror::Error;

/// Errors that can occur during erasure coding operations
#[derive(Debug, Error)]
pub enum ErasureError {
    #[error("Failed to create Reed-Solomon encoder: {0}")]
    CreateError(String),

    #[error("Encoding failed: {0}")]
    EncodeError(String),

    #[error("Decoding failed: {0}")]
    DecodeError(String),

    #[error("Invalid shard count: expected {expected}, got {actual}")]
    InvalidShardCount { expected: usize, actual: usize },

    #[error("Not enough shards for reconstruction: need {needed}, have {available}")]
    NotEnoughShards { needed: usize, available: usize },

    #[error("Shard size mismatch: expected {expected}, got {actual}")]
    ShardSizeMismatch { expected: usize, actual: usize },
}

/// Result type for erasure coding operations
pub type ErasureResult<T> = Result<T, ErasureError>;

/// Reed-Solomon erasure coding codec
///
/// Provides encoding and decoding of data using Reed-Solomon erasure codes.
/// Data is split into `data_shards` pieces, and `parity_shards` parity pieces
/// are generated. Any `data_shards` pieces (data or parity) are sufficient
/// to reconstruct the original data.
#[derive(Debug)]
pub struct Codec {
    rs: ReedSolomon,
    data_shards: usize,
    parity_shards: usize,
}

impl Codec {
    /// Create a new Reed-Solomon codec
    ///
    /// # Arguments
    /// * `data_shards` - Number of data shards (original data is split into this many pieces)
    /// * `parity_shards` - Number of parity shards (fault tolerance)
    ///
    /// # Example
    /// ```
    /// use lightningstor_distributed::erasure::Codec;
    ///
    /// // 4+2 configuration: 4 data shards, 2 parity shards
    /// // Can tolerate loss of any 2 shards
    /// let codec = Codec::new(4, 2).unwrap();
    /// ```
    pub fn new(data_shards: usize, parity_shards: usize) -> ErasureResult<Self> {
        let rs = ReedSolomon::new(data_shards, parity_shards)
            .map_err(|e| ErasureError::CreateError(e.to_string()))?;

        Ok(Self {
            rs,
            data_shards,
            parity_shards,
        })
    }

    /// Get the number of data shards
    pub fn data_shards(&self) -> usize {
        self.data_shards
    }

    /// Get the number of parity shards
    pub fn parity_shards(&self) -> usize {
        self.parity_shards
    }

    /// Get the total number of shards (data + parity)
    pub fn total_shards(&self) -> usize {
        self.data_shards + self.parity_shards
    }

    /// Calculate the shard size for given data size
    ///
    /// Each shard will be this size (data is padded if necessary)
    pub fn shard_size(&self, data_size: usize) -> usize {
        // Round up to ensure all data fits
        (data_size + self.data_shards - 1) / self.data_shards
    }

    /// Encode data into shards
    ///
    /// Returns a vector of shards: first `data_shards` are data shards,
    /// remaining `parity_shards` are parity shards.
    ///
    /// # Arguments
    /// * `data` - The data to encode
    ///
    /// # Returns
    /// A vector of `data_shards + parity_shards` shards, each of equal size.
    pub fn encode(&self, data: &[u8]) -> ErasureResult<Vec<Vec<u8>>> {
        if data.is_empty() {
            // Handle empty data - create minimal shards
            let shard_size = 1;
            let mut shards: Vec<Vec<u8>> = (0..self.total_shards())
                .map(|_| vec![0u8; shard_size])
                .collect();

            self.rs
                .encode(&mut shards)
                .map_err(|e| ErasureError::EncodeError(e.to_string()))?;

            return Ok(shards);
        }

        let shard_size = self.shard_size(data.len());
        let total_shards = self.total_shards();

        // Create shards with padding
        let mut shards: Vec<Vec<u8>> = Vec::with_capacity(total_shards);

        // Fill data shards
        for i in 0..self.data_shards {
            let start = i * shard_size;
            let end = std::cmp::min(start + shard_size, data.len());

            let mut shard = vec![0u8; shard_size];
            if start < data.len() {
                let copy_len = end - start;
                shard[..copy_len].copy_from_slice(&data[start..end]);
            }
            shards.push(shard);
        }

        // Create empty parity shards
        for _ in 0..self.parity_shards {
            shards.push(vec![0u8; shard_size]);
        }

        // Encode (fills in parity shards)
        self.rs
            .encode(&mut shards)
            .map_err(|e| ErasureError::EncodeError(e.to_string()))?;

        Ok(shards)
    }

    /// Decode shards back into original data
    ///
    /// # Arguments
    /// * `shards` - Vector of optional shards. `None` indicates a missing shard.
    /// * `original_size` - The original data size (needed to remove padding)
    ///
    /// # Returns
    /// The reconstructed original data.
    pub fn decode(&self, shards: Vec<Option<Vec<u8>>>, original_size: usize) -> ErasureResult<Vec<u8>> {
        if shards.len() != self.total_shards() {
            return Err(ErasureError::InvalidShardCount {
                expected: self.total_shards(),
                actual: shards.len(),
            });
        }

        // Count available shards
        let available = shards.iter().filter(|s| s.is_some()).count();
        if available < self.data_shards {
            return Err(ErasureError::NotEnoughShards {
                needed: self.data_shards,
                available,
            });
        }

        // Determine shard size from first available shard
        let shard_size = shards
            .iter()
            .find_map(|s| s.as_ref().map(|v| v.len()))
            .unwrap_or(1);

        // Verify all shards have same size
        for shard in shards.iter() {
            if let Some(s) = shard {
                if s.len() != shard_size {
                    return Err(ErasureError::ShardSizeMismatch {
                        expected: shard_size,
                        actual: s.len(),
                    });
                }
            }
        }

        // Convert to the format expected by reed-solomon-erasure
        let mut shard_refs: Vec<Option<Vec<u8>>> = shards;

        // Reconstruct missing shards
        self.rs
            .reconstruct(&mut shard_refs)
            .map_err(|e| ErasureError::DecodeError(e.to_string()))?;

        // Reassemble data from data shards
        let mut result = Vec::with_capacity(original_size);
        for shard in shard_refs.into_iter().take(self.data_shards) {
            if let Some(data) = shard {
                result.extend_from_slice(&data);
            }
        }

        // Trim to original size (remove padding)
        result.truncate(original_size);

        Ok(result)
    }

    /// Verify that shards are consistent
    ///
    /// Returns true if parity shards are correct for the given data shards.
    pub fn verify(&self, shards: &[Vec<u8>]) -> ErasureResult<bool> {
        if shards.len() != self.total_shards() {
            return Err(ErasureError::InvalidShardCount {
                expected: self.total_shards(),
                actual: shards.len(),
            });
        }

        let refs: Vec<&[u8]> = shards.iter().map(|s| s.as_slice()).collect();

        Ok(self.rs.verify(&refs).unwrap_or(false))
    }
}

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

    #[test]
    fn test_codec_creation() {
        let codec = Codec::new(4, 2).unwrap();
        assert_eq!(codec.data_shards(), 4);
        assert_eq!(codec.parity_shards(), 2);
        assert_eq!(codec.total_shards(), 6);
    }

    #[test]
    fn test_encode_decode_roundtrip() {
        let codec = Codec::new(4, 2).unwrap();
        let original = b"Hello, World! This is a test of erasure coding.";

        let shards = codec.encode(original).unwrap();
        assert_eq!(shards.len(), 6);

        // Decode with all shards
        let all_shards: Vec<Option<Vec<u8>>> = shards.into_iter().map(Some).collect();
        let decoded = codec.decode(all_shards, original.len()).unwrap();

        assert_eq!(decoded, original);
    }

    #[test]
    fn test_decode_with_missing_data_shards() {
        let codec = Codec::new(4, 2).unwrap();
        let original = b"Hello, World! This is a test of erasure coding.";

        let shards = codec.encode(original).unwrap();

        // Remove 2 data shards (indices 0 and 1)
        let mut partial_shards: Vec<Option<Vec<u8>>> = shards.into_iter().map(Some).collect();
        partial_shards[0] = None;
        partial_shards[1] = None;

        let decoded = codec.decode(partial_shards, original.len()).unwrap();
        assert_eq!(decoded, original);
    }

    #[test]
    fn test_decode_with_missing_parity_shards() {
        let codec = Codec::new(4, 2).unwrap();
        let original = b"Hello, World! This is a test of erasure coding.";

        let shards = codec.encode(original).unwrap();

        // Remove both parity shards (indices 4 and 5)
        let mut partial_shards: Vec<Option<Vec<u8>>> = shards.into_iter().map(Some).collect();
        partial_shards[4] = None;
        partial_shards[5] = None;

        let decoded = codec.decode(partial_shards, original.len()).unwrap();
        assert_eq!(decoded, original);
    }

    #[test]
    fn test_decode_with_mixed_missing_shards() {
        let codec = Codec::new(4, 2).unwrap();
        let original = b"Hello, World! This is a test of erasure coding.";

        let shards = codec.encode(original).unwrap();

        // Remove 1 data shard and 1 parity shard
        let mut partial_shards: Vec<Option<Vec<u8>>> = shards.into_iter().map(Some).collect();
        partial_shards[2] = None; // data shard
        partial_shards[5] = None; // parity shard

        let decoded = codec.decode(partial_shards, original.len()).unwrap();
        assert_eq!(decoded, original);
    }

    #[test]
    fn test_decode_fails_with_too_many_missing() {
        let codec = Codec::new(4, 2).unwrap();
        let original = b"Hello, World!";

        let shards = codec.encode(original).unwrap();

        // Remove 3 shards (more than parity_shards)
        let mut partial_shards: Vec<Option<Vec<u8>>> = shards.into_iter().map(Some).collect();
        partial_shards[0] = None;
        partial_shards[1] = None;
        partial_shards[2] = None;

        let result = codec.decode(partial_shards, original.len());
        assert!(result.is_err());
    }

    #[test]
    fn test_verify_valid_shards() {
        let codec = Codec::new(4, 2).unwrap();
        let original = b"Test data for verification";

        let shards = codec.encode(original).unwrap();
        assert!(codec.verify(&shards).unwrap());
    }

    #[test]
    fn test_verify_corrupted_shards() {
        let codec = Codec::new(4, 2).unwrap();
        let original = b"Test data for verification";

        let mut shards = codec.encode(original).unwrap();

        // Corrupt a shard
        shards[0][0] ^= 0xFF;

        assert!(!codec.verify(&shards).unwrap());
    }

    #[test]
    fn test_encode_empty_data() {
        let codec = Codec::new(4, 2).unwrap();
        let shards = codec.encode(&[]).unwrap();
        assert_eq!(shards.len(), 6);
    }

    #[test]
    fn test_encode_large_data() {
        let codec = Codec::new(4, 2).unwrap();
        let original: Vec<u8> = (0..10000).map(|i| (i % 256) as u8).collect();

        let shards = codec.encode(&original).unwrap();
        let all_shards: Vec<Option<Vec<u8>>> = shards.into_iter().map(Some).collect();
        let decoded = codec.decode(all_shards, original.len()).unwrap();

        assert_eq!(decoded, original);
    }

    #[test]
    fn test_shard_size_calculation() {
        let codec = Codec::new(4, 2).unwrap();

        // 100 bytes / 4 shards = 25 bytes per shard
        assert_eq!(codec.shard_size(100), 25);

        // 101 bytes / 4 shards = 26 bytes per shard (rounded up)
        assert_eq!(codec.shard_size(101), 26);

        // 0 bytes = 0 bytes per shard
        assert_eq!(codec.shard_size(0), 0);
    }
}
