self_encryption 0.35.0

// Copyright 2025 MaidSafe.net limited.
//
// This SAFE Network Software is licensed to you under The General Public License (GPL), version 3.
// Unless required by applicable law or agreed to in writing, the SAFE Network Software distributed
// under the GPL Licence is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. Please review the Licences for the specific language governing
// permissions and limitations relating to use of the SAFE Network Software.

//! Streaming encryption functionality for memory-efficient processing of large files.

use crate::{
    encrypt::encrypt_chunk,
    shrink_data_map,
    utils::{get_num_chunks, get_pad_key_and_nonce, get_start_end_positions},
    ChunkInfo, DataMap, EncryptedChunk, Error, Result,
};
use bytes::Bytes;
use xor_name::XorName;

/// Internal return type for the streaming encryption processing.
#[derive(Debug, Clone)]
enum ChunkOrDataMap {
    /// An encrypted chunk ready for storage
    Chunk(EncryptedChunk),
    /// The final DataMap when encryption is complete
    DataMap(DataMap),
}

/// Iterator that yields only encrypted chunks from an EncryptionStream.
///
/// This provides a clean interface for iterating over chunks as (XorName, Bytes) tuples.
/// The DataMap becomes available on the parent EncryptionStream after iteration completes.
pub struct ChunkStream<'a, I> {
    stream: &'a mut EncryptionStream<I>,
}

impl<'a, I> Iterator for ChunkStream<'a, I>
where
    I: Iterator<Item = Bytes>,
{
    type Item = Result<(XorName, Bytes)>;

    fn next(&mut self) -> Option<Self::Item> {
        // Get the next item from the underlying stream
        match self.stream.next_internal() {
            Some(Ok(ChunkOrDataMap::Chunk(chunk))) => {
                // Convert to (hash, content) tuple for direct usage
                let hash = crate::hash::content_hash(&chunk.content);
                let content = chunk.content;
                Some(Ok((hash, content)))
            }
            Some(Ok(ChunkOrDataMap::DataMap(datamap))) => {
                // Store the datamap - this signals we're done with chunks
                self.stream.final_datamap = Some(datamap);
                None
            }
            Some(Err(e)) => Some(Err(e)),
            None => None,
        }
    }
}

/// Streaming encryption processor that handles data and produces encrypted chunks.
///
/// This processes incoming data bytes, splits them into chunks according to
/// self-encryption rules, encrypts them, and makes them available through iteration.
/// At the end, it applies shrinking and provides access to the final DataMap.
///
/// # Memory Efficiency
///
/// This implementation maintains minimal memory usage by:
/// - Only buffering incomplete chunks (typically < 1MB)  
/// - Processing data sequentially as it arrives
/// - Applying shrinking at the end
pub struct EncryptionStream<I> {
    /// Input data iterator
    data_iter: I,
    /// Total expected data size
    data_size: usize,
    /// Buffer for accumulating data into complete chunks
    buffer: Vec<u8>,
    /// Number of chunks processed so far
    chunks_processed: usize,
    /// Total number of chunks expected
    total_chunks: usize,
    /// Source hashes of all chunks (filled as we process them)
    src_hashes: Vec<Option<XorName>>,
    /// Raw data for deferred chunks 0 and 1
    deferred_chunks: [Option<Bytes>; 2],
    /// Chunk infos for building the final DataMap
    chunk_infos: Vec<Option<ChunkInfo>>,
    /// Whether we've finished processing all input data
    input_complete: bool,
    /// Whether we've yielded the final DataMap
    is_complete: bool,
    /// Final DataMap (available after chunks iteration completes)
    final_datamap: Option<DataMap>,
    /// Shrinking chunks to yield after main processing
    shrinking_chunks: Vec<EncryptedChunk>,
    /// Current index for yielding shrinking chunks
    shrinking_chunk_index: usize,
}

impl<I> EncryptionStream<I>
where
    I: Iterator<Item = Bytes>,
{
    /// Creates a new streaming encryption iterator.
    ///
    /// # Arguments
    ///
    /// * `data_size` - The total size of the data to be encrypted
    /// * `data_iter` - Iterator providing data chunks as `Bytes`
    ///
    /// # Returns
    ///
    /// * `Result<EncryptionStream<I>>` - A new streaming encryption iterator
    ///
    /// # Examples
    ///
    /// ```rust
    /// use self_encryption::stream_encrypt;
    /// use bytes::Bytes;
    ///
    /// # fn main() -> self_encryption::Result<()> {
    /// let file_data = b"Hello, world! This is test data for streaming encryption.".repeat(1000);
    /// let data_size = file_data.len();
    /// let data_iter = file_data.chunks(1024).map(|chunk| Bytes::from(chunk.to_vec()));
    ///
    /// let mut stream = stream_encrypt(data_size, data_iter)?;
    /// let mut encrypted_chunks = Vec::new();
    ///
    /// // Use the clean chunks() API
    /// for chunk_result in stream.chunks() {
    ///     let (_hash, _content) = chunk_result?;
    ///     encrypted_chunks.push((_hash, _content));
    /// }
    ///
    /// let _data_map = stream.datamap().expect("Should have DataMap");
    /// assert!(encrypted_chunks.len() > 0);
    /// # Ok(())
    /// # }
    /// ```
    fn new(data_size: usize, data_iter: I) -> Result<Self> {
        if data_size < crate::MIN_ENCRYPTABLE_BYTES {
            return Err(Error::Generic(format!(
                "File too small for self-encryption! Required size at least {}",
                crate::MIN_ENCRYPTABLE_BYTES
            )));
        }

        let total_chunks = get_num_chunks(data_size);
        if total_chunks < 3 {
            return Err(Error::Generic(
                "File must be large enough to generate at least 3 chunks".to_string(),
            ));
        }

        Ok(Self {
            data_iter,
            data_size,
            buffer: Vec::new(),
            chunks_processed: 0,
            total_chunks,
            src_hashes: vec![None; total_chunks],
            deferred_chunks: [None, None],
            chunk_infos: vec![None; total_chunks],
            input_complete: false,
            is_complete: false,
            final_datamap: None,
            shrinking_chunks: Vec::new(),
            shrinking_chunk_index: 0,
        })
    }

    /// Process any complete chunks from the buffer and yield them
    fn try_process_chunks(&mut self) -> Result<Option<ChunkOrDataMap>> {
        while self.chunks_processed < self.total_chunks {
            let chunk_index = self.chunks_processed;
            let (chunk_start, chunk_end) = get_start_end_positions(self.data_size, chunk_index);
            let chunk_size = chunk_end - chunk_start;

            // Check if we have enough data for this chunk
            let buffer_start = chunk_start.saturating_sub(self.get_processed_bytes());

            if self.buffer.len() >= buffer_start + chunk_size || self.input_complete {
                // We can process this chunk
                let actual_chunk_size = if self.input_complete {
                    std::cmp::min(chunk_size, self.buffer.len().saturating_sub(buffer_start))
                } else {
                    chunk_size
                };

                if actual_chunk_size == 0 {
                    break;
                }

                let chunk_data = if buffer_start + actual_chunk_size <= self.buffer.len() {
                    Bytes::from(
                        self.buffer[buffer_start..buffer_start + actual_chunk_size].to_vec(),
                    )
                } else {
                    break; // Not enough data yet
                };

                // Calculate source hash and store chunk info
                let src_hash = crate::hash::content_hash(&chunk_data);
                let entry = self.src_hashes.get_mut(chunk_index).ok_or_else(|| {
                    Error::Generic(format!(
                        "chunk index {chunk_index} out of bounds for src_hashes"
                    ))
                })?;
                *entry = Some(src_hash);
                let info_entry = self.chunk_infos.get_mut(chunk_index).ok_or_else(|| {
                    Error::Generic(format!(
                        "chunk index {chunk_index} out of bounds for chunk_infos"
                    ))
                })?;
                *info_entry = Some(ChunkInfo {
                    index: chunk_index,
                    dst_hash: crate::hash::content_hash(&[]), // Will be filled when encrypted
                    src_hash,
                    src_size: chunk_data.len(),
                });

                // Handle encryption based on chunk index
                if chunk_index < 2 {
                    // Defer first two chunks until we have all source hashes
                    let deferred_entry =
                        self.deferred_chunks.get_mut(chunk_index).ok_or_else(|| {
                            Error::Generic(format!(
                                "chunk index {chunk_index} out of bounds for deferred_chunks"
                            ))
                        })?;
                    *deferred_entry = Some(chunk_data);
                } else if self.can_encrypt_chunk(chunk_index) {
                    // Encrypt and yield chunk immediately (streaming behavior)
                    let encrypted_chunk = self.encrypt_chunk(chunk_index, chunk_data)?;
                    self.chunks_processed += 1;

                    // Remove processed data from buffer
                    if buffer_start + actual_chunk_size <= self.buffer.len() {
                        let _ = self
                            .buffer
                            .drain(buffer_start..buffer_start + actual_chunk_size);
                    }

                    return Ok(Some(ChunkOrDataMap::Chunk(encrypted_chunk)));
                }

                self.chunks_processed += 1;

                // Remove processed data from buffer
                if buffer_start + actual_chunk_size <= self.buffer.len() {
                    let _ = self
                        .buffer
                        .drain(buffer_start..buffer_start + actual_chunk_size);
                }
            } else {
                // Need more data for this chunk
                break;
            }
        }

        // Check if we can finalize encryption
        if self.input_complete && self.chunks_processed >= self.total_chunks {
            return self.finalize_encryption();
        }

        Ok(None)
    }

    /// Calculate how many bytes we've processed so far
    fn get_processed_bytes(&self) -> usize {
        let mut processed = 0;
        for i in 0..self.chunks_processed {
            let (start, end) = get_start_end_positions(self.data_size, i);
            processed += end - start;
        }
        processed
    }

    /// Check if a chunk can be encrypted (has required source hashes)
    fn can_encrypt_chunk(&self, chunk_index: usize) -> bool {
        if chunk_index < 2 {
            // First two chunks need ALL source hashes
            self.src_hashes.iter().all(|h| h.is_some())
        } else {
            // Chunks 2+ need their own hash and the two dependencies
            let Ok((n1, n2)) = crate::utils::get_n_1_n_2(chunk_index, self.total_chunks) else {
                return false;
            };
            self.src_hashes
                .get(chunk_index)
                .and_then(|h| h.as_ref())
                .is_some()
                && self.src_hashes.get(n1).and_then(|h| h.as_ref()).is_some()
                && self.src_hashes.get(n2).and_then(|h| h.as_ref()).is_some()
        }
    }

    /// Encrypt a chunk and return the indexed encrypted chunk
    fn encrypt_chunk(&mut self, chunk_index: usize, chunk_data: Bytes) -> Result<EncryptedChunk> {
        // Get source hashes for encryption
        let mut src_hashes = vec![crate::hash::content_hash(&[]); self.total_chunks];
        for (i, hash_opt) in self.src_hashes.iter().enumerate() {
            if let Some(hash) = hash_opt {
                let entry = src_hashes.get_mut(i).ok_or_else(|| {
                    Error::Generic(format!("index {i} out of bounds for src_hashes"))
                })?;
                *entry = *hash;
            }
        }

        // Encrypt the chunk (initial encryption is always child_level 0)
        let pki = get_pad_key_and_nonce(chunk_index, &src_hashes, 0)?;
        let encrypted_content = encrypt_chunk(chunk_data, pki)?;
        let dst_hash = crate::hash::content_hash(&encrypted_content);

        // Update chunk info with destination hash
        if let Some(chunk_info) = self
            .chunk_infos
            .get_mut(chunk_index)
            .and_then(|c| c.as_mut())
        {
            chunk_info.dst_hash = dst_hash;
        }

        Ok(EncryptedChunk {
            content: encrypted_content,
        })
    }

    /// Finalize encryption and return the DataMap
    fn finalize_encryption(&mut self) -> Result<Option<ChunkOrDataMap>> {
        if self.is_complete {
            return Ok(None);
        }

        // Collect all source hashes
        let src_hashes: Result<Vec<XorName>> = self
            .src_hashes
            .iter()
            .enumerate()
            .map(|(i, h)| {
                h.ok_or_else(|| Error::Generic(format!("Missing source hash for chunk {i}")))
            })
            .collect();
        let src_hashes = src_hashes?;

        // Process any deferred chunks that haven't been encrypted yet
        for chunk_index in 0..2.min(self.total_chunks) {
            let chunk_data = self
                .deferred_chunks
                .get_mut(chunk_index)
                .and_then(|entry| entry.take());
            if let Some(chunk_data) = chunk_data {
                // For deferred chunks, we need to encrypt them with full source hashes
                let pki = get_pad_key_and_nonce(chunk_index, &src_hashes, 0)?;
                let encrypted_content = encrypt_chunk(chunk_data, pki)?;
                let dst_hash = crate::hash::content_hash(&encrypted_content);

                // Update chunk info
                if let Some(chunk_info) = self
                    .chunk_infos
                    .get_mut(chunk_index)
                    .and_then(|c| c.as_mut())
                {
                    chunk_info.dst_hash = dst_hash;
                }

                // Yield the deferred chunk immediately
                return Ok(Some(ChunkOrDataMap::Chunk(EncryptedChunk {
                    content: encrypted_content,
                })));
            }
        }

        // Build final DataMap
        let mut final_chunk_infos = Vec::new();
        for chunk_info_opt in &self.chunk_infos {
            if let Some(chunk_info) = chunk_info_opt {
                final_chunk_infos.push(chunk_info.clone());
            } else {
                return Err(Error::Generic("Missing chunk info".to_string()));
            }
        }

        final_chunk_infos.sort_by_key(|info| info.index);
        let data_map = DataMap::new(final_chunk_infos);

        // Apply shrinking like streaming_encrypt_from_file does
        let (shrunk_map, shrink_chunks) = shrink_data_map(data_map, |_hash, _content| {
            // Ignore the functor - we just want the returned chunks
            Ok(())
        })?;

        // Store the chunks returned by shrink_data_map for yielding
        self.shrinking_chunks = shrink_chunks;

        // Store the shrunk datamap
        self.final_datamap = Some(shrunk_map);

        // If we have shrinking chunks, yield the first one
        if !self.shrinking_chunks.is_empty()
            && self.shrinking_chunk_index < self.shrinking_chunks.len()
        {
            let chunk = self
                .shrinking_chunks
                .get(self.shrinking_chunk_index)
                .ok_or_else(|| Error::Generic("Shrinking chunk index out of bounds".to_string()))?
                .clone();
            self.shrinking_chunk_index += 1;

            return Ok(Some(ChunkOrDataMap::Chunk(chunk)));
        }

        // No more shrinking chunks, mark as complete and return DataMap
        self.is_complete = true;
        match self.final_datamap.as_ref() {
            Some(dm) => Ok(Some(ChunkOrDataMap::DataMap(dm.clone()))),
            None => Err(Error::Generic(
                "Final datamap missing after finalization".to_string(),
            )),
        }
    }
}

impl<I> EncryptionStream<I>
where
    I: Iterator<Item = Bytes>,
{
    /// Internal method to get the next item from the stream processing (used by ChunkStream)
    fn next_internal(&mut self) -> Option<Result<ChunkOrDataMap>> {
        loop {
            // Check if we have shrinking chunks to yield after processing is complete
            if !self.shrinking_chunks.is_empty()
                && self.shrinking_chunk_index < self.shrinking_chunks.len()
            {
                let chunk = match self.shrinking_chunks.get(self.shrinking_chunk_index) {
                    Some(c) => c.clone(),
                    None => {
                        return Some(Err(Error::Generic(
                            "Shrinking chunk index out of bounds".to_string(),
                        )));
                    }
                };
                self.shrinking_chunk_index += 1;

                // If we've yielded all shrinking chunks, we can yield the final DataMap next
                if self.shrinking_chunk_index >= self.shrinking_chunks.len() {
                    self.is_complete = true;
                }

                return Some(Ok(ChunkOrDataMap::Chunk(chunk)));
            }

            if self.is_complete {
                // Check if we have a final datamap to yield
                if let Some(datamap) = &self.final_datamap {
                    return Some(Ok(ChunkOrDataMap::DataMap(datamap.clone())));
                }

                return None;
            }

            // First, try to process any complete chunks from existing buffer
            match self.try_process_chunks() {
                Ok(Some(result)) => return Some(Ok(result)),
                Ok(None) => {} // No chunks ready, need more data
                Err(e) => return Some(Err(e)),
            }

            // If no chunks are ready, try to get more data
            if !self.input_complete {
                match self.data_iter.next() {
                    Some(data) => {
                        self.buffer.extend_from_slice(&data);

                        // Try processing again with new data
                        match self.try_process_chunks() {
                            Ok(Some(result)) => return Some(Ok(result)),
                            Ok(None) => continue, // Loop to try getting more data
                            Err(e) => return Some(Err(e)),
                        }
                    }
                    None => {
                        // No more input data
                        self.input_complete = true;

                        // Try final processing
                        match self.try_process_chunks() {
                            Ok(Some(result)) => return Some(Ok(result)),
                            Ok(None) => return None, // All done
                            Err(e) => return Some(Err(e)),
                        }
                    }
                }
            } else {
                // Input is complete but no more results
                return None;
            }
        }
    }
}

impl<I> EncryptionStream<I> {
    /// Returns an iterator that yields only encrypted chunks.
    ///
    /// This provides a clean interface for processing chunks without having to
    /// pattern match on the ChunkOrDataMap enum. After the returned iterator
    /// is consumed, the DataMap will be available via `datamap()`.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use self_encryption::stream_encrypt;
    /// use bytes::Bytes;
    ///
    /// # fn main() -> self_encryption::Result<()> {
    /// let file_data = b"Hello, world!".repeat(1000);
    /// let data_iter = file_data.chunks(1024).map(|chunk| Bytes::from(chunk.to_vec()));
    ///
    /// let mut stream = stream_encrypt(file_data.len(), data_iter)?;
    ///
    /// // Clean iteration over chunks only!
    /// for chunk_result in stream.chunks() {
    ///     let (_hash, _content) = chunk_result?;
    ///     // println!("Got chunk {} with {} bytes", hex::encode(hash), content.len());
    ///     // Store chunk directly to your backend
    ///     // store(hash, content)?;
    /// }
    ///
    /// // Get the final DataMap
    /// let _datamap = stream.datamap().expect("Should have DataMap after chunks iteration");
    /// # Ok(())
    /// # }
    /// ```
    pub fn chunks(&mut self) -> ChunkStream<'_, I> {
        ChunkStream { stream: self }
    }

    /// Returns the final DataMap after chunk iteration is complete.
    ///
    /// This method should be called after the `chunks()` iterator has been
    /// fully consumed. Returns `None` if encryption is not yet complete.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use self_encryption::stream_encrypt;
    /// use bytes::Bytes;
    ///
    /// # fn main() -> self_encryption::Result<()> {
    /// let file_data = b"Hello, world!".repeat(1000);
    /// let data_iter = file_data.chunks(1024).map(|chunk| Bytes::from(chunk.to_vec()));
    ///
    /// let mut stream = stream_encrypt(file_data.len(), data_iter)?;
    ///
    /// // Process all chunks
    /// for chunk_result in stream.chunks() {
    ///     let (_hash, _content) = chunk_result?;
    ///     // Store chunk directly
    ///     // store(hash, content)?;
    /// }
    ///
    /// // Get the DataMap
    /// let _datamap = stream.datamap().expect("Should have DataMap");
    /// # Ok(())
    /// # }
    /// ```
    pub fn datamap(&self) -> Option<&DataMap> {
        self.final_datamap.as_ref()
    }

    /// Returns the final DataMap after chunk iteration is complete, consuming the stream.
    ///
    /// This method should be called after the `chunks()` iterator has been
    /// fully consumed. Returns `None` if encryption is not yet complete.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use self_encryption::stream_encrypt;
    /// use bytes::Bytes;
    ///
    /// # fn main() -> self_encryption::Result<()> {
    /// let file_data = b"Hello, world!".repeat(1000);
    /// let data_iter = file_data.chunks(1024).map(|chunk| Bytes::from(chunk.to_vec()));
    ///
    /// let mut stream = stream_encrypt(file_data.len(), data_iter)?;
    ///
    /// // Process all chunks
    /// for chunk_result in stream.chunks() {
    ///     let (_hash, _content) = chunk_result?;
    ///     // Store chunk directly
    ///     // store(hash, content)?;
    /// }
    ///
    /// // Get the DataMap (consuming the stream)
    /// let _datamap = stream.into_datamap();
    /// # Ok(())
    /// # }
    /// ```
    /// Returns the final DataMap, or `None` if encryption is not yet complete
    /// (i.e., `chunks()` was not fully consumed).
    pub fn into_datamap(self) -> Option<DataMap> {
        self.final_datamap
    }
}

/// Creates a streaming encryption iterator that processes data on-the-fly.
///
/// This function provides memory-efficient encryption by processing data as it flows
/// through an iterator, yielding encrypted chunks immediately without buffering them.
/// Only a small amount of data is kept in memory for incomplete chunks.
///
/// # Arguments
///
/// * `data_size` - The total size of the data to be encrypted
/// * `data_iter` - An iterator providing data chunks as `Bytes`
///
/// # Returns
///
/// * `Result<EncryptionStream<I>>` - A streaming encryption iterator
///
/// # Examples
///
/// ## Basic Usage
///
/// ```rust
/// use self_encryption::stream_encrypt;
/// use bytes::Bytes;
///
/// # fn main() -> self_encryption::Result<()> {
/// // Create test data
/// let file_data = b"Hello, world! ".repeat(5000); // ~65KB
/// let data_size = file_data.len();
///
/// // Create iterator over data chunks
/// let data_iter = file_data.chunks(1024).map(|chunk| Bytes::from(chunk.to_vec()));
///
/// // Stream encrypt the data - now much cleaner!
/// let mut stream = stream_encrypt(data_size, data_iter)?;
///
/// for chunk_result in stream.chunks() {
///     let (hash, content) = chunk_result?;
///     println!("Got encrypted chunk {} with {} bytes", hex::encode(hash), content.len());
///     // Store chunk directly to your backend
///     // store(hash, content)?;
/// }
///
/// let data_map = stream.datamap().expect("Should have DataMap after iteration");
/// println!("Encryption complete! DataMap has {} chunks", data_map.len());
/// # Ok(())
/// # }
/// ```
///
/// ## File Processing
///
/// ```rust
/// use self_encryption::stream_encrypt;
/// use bytes::Bytes;
/// use std::fs::File;
/// use std::io::{BufReader, Read};
///
/// # fn main() -> std::result::Result<(), Box<dyn std::error::Error>> {
/// # use tempfile::NamedTempFile;
/// # use std::io::Write;
/// # let mut temp_file = NamedTempFile::new()?;
/// # temp_file.write_all(&vec![42u8; 100000])?;
/// # let file_path = temp_file.path();
///
/// let file = File::open(file_path)?;
/// let data_size = file.metadata()?.len() as usize;
/// let mut reader = BufReader::new(file);
///
/// // Create iterator that reads file in chunks
/// let data_iter = std::iter::from_fn(move || {
///     let mut buffer = vec![0u8; 8192];
///     match reader.read(&mut buffer) {
///         Ok(0) => None, // EOF
///         Ok(n) => {
///             buffer.truncate(n);
///             Some(Bytes::from(buffer))
///         }
///         Err(_) => None,
///     }
/// });
///
/// // Process the file - much cleaner!
/// let mut stream = stream_encrypt(data_size, data_iter)?;
/// for chunk_result in stream.chunks() {
///     let (hash, content) = chunk_result?;
///     // Store chunk to your preferred backend
///     println!("Storing chunk {} ({} bytes)", hex::encode(hash), content.len());
/// }
/// println!("File encrypted successfully!");
/// let _datamap = stream.datamap().expect("Should have DataMap");
/// # Ok(())
/// # }
/// ```
pub fn stream_encrypt<I>(data_size: usize, data_iter: I) -> Result<EncryptionStream<I>>
where
    I: Iterator<Item = Bytes>,
{
    EncryptionStream::new(data_size, data_iter)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_helpers::random_bytes;
    use std::collections::HashMap;

    #[test]
    fn test_stream_encrypt_basic() -> Result<()> {
        let data_size = 50_000; // 50KB
        let original_data = random_bytes(data_size);
        let data_iter = original_data
            .chunks(1024)
            .map(|chunk| Bytes::from(chunk.to_vec()));

        let mut encrypted_chunks = Vec::new();
        let mut stream = stream_encrypt(data_size, data_iter)?;

        // Clean iteration with new API!
        for chunk_result in stream.chunks() {
            let (hash, content) = chunk_result?;
            encrypted_chunks.push((hash, content));
        }

        let data_map = stream
            .datamap()
            .expect("Should have completed with DataMap");
        assert_ne!(data_map.len(), 0, "DataMap should have chunks");
        assert_ne!(
            encrypted_chunks.len(),
            0,
            "Should have yielded encrypted chunks"
        );

        Ok(())
    }

    #[test]
    fn test_stream_encrypt_single_chunk() -> Result<()> {
        let data_size = 10_000; // Small file
        let original_data = random_bytes(data_size);
        let data_iter = std::iter::once(original_data);

        let mut chunks = Vec::new();
        let mut stream = stream_encrypt(data_size, data_iter)?;

        for chunk_result in stream.chunks() {
            let (hash, content) = chunk_result?;
            chunks.push((hash, content));
        }

        // Should get chunks and DataMap should be available
        assert!(!chunks.is_empty());
        let _datamap = stream.datamap().expect("Should have DataMap");

        Ok(())
    }

    #[test]
    fn test_stream_encrypt_large_file() -> Result<()> {
        let data_size = 5_000_000; // 5MB
        let original_data = random_bytes(data_size);
        let chunk_size = 64 * 1024; // 64KB chunks
        let data_iter = original_data
            .chunks(chunk_size)
            .map(|chunk| Bytes::from(chunk.to_vec()));

        let mut encrypted_chunks = Vec::new();
        let mut stream = stream_encrypt(data_size, data_iter)?;

        for chunk_result in stream.chunks() {
            let (hash, content) = chunk_result?;
            encrypted_chunks.push((hash, content));
        }

        let _data_map = stream.datamap().expect("Should complete with DataMap");
        assert!(
            encrypted_chunks.len() > 1,
            "Large file should produce multiple chunks"
        );

        Ok(())
    }

    #[test]
    fn test_stream_encrypt_memory_efficiency() -> Result<()> {
        let data_size = 10_000_000; // 10MB
        let chunk_size = 8192; // Small input chunks to test buffering

        // Create iterator that tracks memory usage
        let original_data = random_bytes(data_size);
        let data_iter = original_data
            .chunks(chunk_size)
            .map(|chunk| Bytes::from(chunk.to_vec()));

        let mut stream = stream_encrypt(data_size, data_iter)?;
        let _max_buffer_size = 0;
        let mut encrypted_count = 0;

        // Process stream and monitor buffer size
        for chunk_result in stream.chunks() {
            let (_hash, _content) = chunk_result?;
            encrypted_count += 1;
            // In a real implementation, we'd check buffer size here
            // For now, just ensure we're yielding chunks progressively
        }

        assert!(
            encrypted_count > 0,
            "Should yield encrypted chunks progressively"
        );
        Ok(())
    }

    #[test]
    fn test_stream_encrypt_too_small() {
        let data_size = 2; // Too small (less than MIN_ENCRYPTABLE_BYTES = 3)
        let data = vec![42u8; data_size];
        let data_iter = std::iter::once(Bytes::from(data));

        let result = stream_encrypt(data_size, data_iter);
        assert!(
            result.is_err(),
            "Should reject files that are too small for self-encryption"
        );
    }

    #[test]
    fn test_stream_encrypt_consistency() -> Result<()> {
        // Test that streaming encryption produces same result as standard encryption
        let data_size = 200_000;
        let original_data = random_bytes(data_size);

        // Standard encryption
        let (_standard_data_map, standard_chunks) = crate::encrypt(original_data.clone())?;

        // Streaming encryption
        let data_iter = original_data
            .chunks(4096)
            .map(|chunk| Bytes::from(chunk.to_vec()));
        let mut streaming_chunks = Vec::new();
        let mut stream = stream_encrypt(data_size, data_iter)?;

        for chunk_result in stream.chunks() {
            let (hash, content) = chunk_result?;
            streaming_chunks.push((hash, content));
        }

        let _streaming_data_map = stream.datamap().expect("Should have DataMap");

        // Both should decrypt to same original data
        let mut standard_storage = HashMap::new();
        for chunk in standard_chunks {
            let hash = crate::hash::content_hash(&chunk.content);
            let _ = standard_storage.insert(hash, chunk.content.to_vec());
        }

        let mut streaming_storage = HashMap::new();
        for (hash, content) in streaming_chunks {
            let _ = streaming_storage.insert(hash, content.to_vec());
        }

        // Both methods now apply shrinking, so they should produce similar chunk counts
        assert_ne!(standard_storage.len(), 0, "Standard should produce chunks");
        assert_ne!(
            streaming_storage.len(),
            0,
            "Streaming should produce chunks"
        );

        Ok(())
    }
}