keep/src/compression_engine.rs

use anyhow::{Result, anyhow};
use std::io;
use std::io::{Read, Write};
use std::path::PathBuf;
use strum::IntoEnumIterator;

use log::*;

use lazy_static::lazy_static;

extern crate enum_map;
use enum_map::enum_map;
use enum_map::{Enum, EnumMap};

pub mod gzip;
pub mod lz4;
pub mod none;
pub mod program;

use crate::compression_engine::gzip::CompressionEngineGZip;
use crate::compression_engine::lz4::CompressionEngineLZ4;
use crate::compression_engine::none::CompressionEngineNone;
use crate::compression_engine::program::CompressionEngineProgram;

/// Enum representing different compression types supported by the system
#[derive(Debug, Eq, PartialEq, Clone, strum::EnumIter, strum::Display, strum::EnumString, Enum)]
#[strum(ascii_case_insensitive)]
pub enum CompressionType {
    LZ4,
    GZip,
    BZip2,
    XZ,
    ZStd,
    None,
}

/// Trait defining the interface for compression engines
pub trait CompressionEngine {
    /// Opens a compressed file for reading
    ///
    /// # Arguments
    ///
    /// * `file_path` - Path to the compressed file
    ///
    /// # Returns
    ///
    /// * `Result<Box<dyn Read>>` - A boxed reader that decompresses the file on read
    fn open(&self, file_path: PathBuf) -> Result<Box<dyn Read>>;
    
    /// Creates a new compressed file for writing
    ///
    /// # Arguments
    ///
    /// * `file_path` - Path where the compressed file will be created
    ///
    /// # Returns
    ///
    /// * `Result<Box<dyn Write>>` - A boxed writer that compresses data on write
    fn create(&self, file_path: PathBuf) -> Result<Box<dyn Write>>;

    /// Checks if this compression engine is supported on the current system
    ///
    /// # Returns
    ///
    /// * `bool` - True if supported, false otherwise
    fn is_supported(&self) -> bool {
        true
    }

    /// Copies decompressed content from a file to a writer
    ///
    /// # Arguments
    ///
    /// * `file_path` - Path to the compressed file
    /// * `writer` - Writer to receive decompressed content
    ///
    /// # Returns
    ///
    /// * `Result<()>` - Success or error
    fn copy(&self, file_path: PathBuf, writer: &mut dyn Write) -> Result<()> {
        let mut reader = self.open(file_path)?;
        io::copy(&mut reader, writer)?;
        writer.flush()?;
        Ok(())
    }

    /// Decompresses and outputs file content to stdout
    ///
    /// # Arguments
    ///
    /// * `file_path` - Path to the compressed file
    ///
    /// # Returns
    ///
    /// * `Result<()>` - Success or error
    fn cat(&self, file_path: PathBuf) -> Result<()> {
        let mut stdout = io::stdout().lock();
        self.copy(file_path, &mut stdout)
    }

    /// Calculates the decompressed size of a file
    ///
    /// # Arguments
    ///
    /// * `file_path` - Path to the compressed file
    ///
    /// # Returns
    ///
    /// * `Result<usize>` - The decompressed size in bytes
    fn size(&self, file_path: PathBuf) -> Result<usize> {
        let mut reader = self.open(file_path)?;
        let mut buffer = [0; libc::BUFSIZ as usize];
        let mut size: usize = 0;

        loop {
            let n = reader.read(&mut buffer[..libc::BUFSIZ as usize])?;
            if n == 0 {
                debug!("COMPRESSION: EOF");
                break;
            }

            size += n;
        }

        Ok(size)
    }
}

lazy_static! {
    /// Mapping of compression types to their external program implementations
    pub static ref COMPRESSION_PROGRAMS: EnumMap<CompressionType, Option<CompressionEngineProgram>> = enum_map! {
        CompressionType::LZ4 => {
            let program = CompressionEngineProgram::new("lz4", vec!["-c"], vec!["-d", "-c"]);
            if program.supported { Some(program) } else { None }
        },
        CompressionType::GZip => {
            let program = CompressionEngineProgram::new("gzip", vec!["-c"], vec!["-d", "-c"]);
            if program.supported { Some(program) } else { None }
        },
        #[cfg(feature = "bzip2")]
        CompressionType::BZip2 => {
            let program = CompressionEngineProgram::new("bzip2", vec!["-qcf"], vec!["-dcf"]);
            if program.supported { Some(program) } else { None }
        },
        #[cfg(not(feature = "bzip2"))]
        CompressionType::BZip2 => {
            let program = CompressionEngineProgram::new("bzip2", vec!["-qcf"], vec!["-dcf"]);
            if program.supported { Some(program) } else { None }
        },
        CompressionType::XZ => {
            let program = CompressionEngineProgram::new("xz", vec!["-qcf"], vec!["-dcf"]);
            if program.supported { Some(program) } else { None }
        },
        CompressionType::ZStd => {
            let program = CompressionEngineProgram::new("zstd", vec!["-qcf"], vec!["-dcf"]);
            if program.supported { Some(program) } else { None }
        },
        CompressionType::None => None
    };
}

/// Gets a compression engine for the specified compression type
///
/// # Arguments
///
/// * `compression_type` - The type of compression to use
///
/// # Returns
///
/// * `Result<Box<dyn CompressionEngine>>` - A boxed compression engine instance
pub fn get_compression_engine(
    compression_type: CompressionType,
) -> Result<Box<dyn CompressionEngine>> {
    match compression_type {
        CompressionType::LZ4 => {
            #[cfg(feature = "lz4")]
            {
                Ok(Box::new(CompressionEngineLZ4::new()))
            }
            #[cfg(not(feature = "lz4"))]
            {
                if let Some(engine) = COMPRESSION_PROGRAMS[compression_type].clone() {
                    Ok(Box::new(engine))
                } else {
                    Err(anyhow!("LZ4 not available: feature disabled and program not found"))
                }
            }
        },
        CompressionType::GZip => {
            #[cfg(feature = "gzip")]
            {
                Ok(Box::new(CompressionEngineGZip::new()))
            }
            #[cfg(not(feature = "gzip"))]
            {
                if let Some(engine) = COMPRESSION_PROGRAMS[compression_type].clone() {
                    Ok(Box::new(engine))
                } else {
                    Err(anyhow!("GZip not available: feature disabled and program not found"))
                }
            }
        },
        CompressionType::None => Ok(Box::new(CompressionEngineNone::new())),
        compression_type => {
            let ct = compression_type.clone();
            if let Some(engine) = COMPRESSION_PROGRAMS[ct.clone()].clone() {
                Ok(Box::new(engine))
            } else {
                Err(anyhow!("Compression type {:?} not supported", ct))
            }
        },
    }
}

/// Gets the default compression type based on available support
///
/// # Returns
///
/// * `CompressionType` - The first supported compression type found
pub fn default_compression_type() -> CompressionType {
    let mut default = CompressionType::None;
    for compression_type in CompressionType::iter() {
        let compression_engine =
            get_compression_engine(compression_type.clone()).expect("Missing engine");
        if compression_engine.is_supported() {
            default = compression_type;
            break;
        }
    }
    default
}