s3m-core 0.21.0

use crate::s3::error::Error as S3Error;
use crate::s3::limits::{MAX_PART_SIZE_BYTES, MAX_PARTS_PER_UPLOAD};
use base64ct::{Base64, Encoding};
use ring::{digest, hmac};
use std::{
    fmt::Write,
    io::{Error, ErrorKind, Read},
    path::Path,
};
use tokio::time::{Duration, sleep};

/// Recommended default number of concurrent upload/download requests.
///
/// Uses [`std::thread::available_parallelism`], which respects cgroup CPU
/// quotas and CPU affinity, unlike a raw physical-core count. This matters for
/// the constrained / containerized environments `s3m` is designed to run in: a
/// container limited to 2 CPUs reports `2` here even on a 64-core host. Two
/// units are reserved for the async runtime and the main task, and the result
/// is always at least `1`.
#[must_use]
pub fn default_concurrency() -> usize {
    std::thread::available_parallelism()
        .map_or(1, std::num::NonZeroUsize::get)
        .saturating_sub(2)
        .max(1)
}

#[must_use]
pub fn sha256_digest(input: impl AsRef<[u8]>) -> digest::Digest {
    digest::digest(&digest::SHA256, input.as_ref())
}

#[must_use]
pub fn base64_md5(input: impl AsRef<[u8]>) -> String {
    let md5_digest = md5::compute(input);
    Base64::encode_string(md5_digest.as_ref())
}

#[must_use]
pub fn sha256_hmac(key: &[u8], msg: &[u8]) -> hmac::Tag {
    let s_key = hmac::Key::new(hmac::HMAC_SHA256, key);
    hmac::sign(&s_key, msg)
}

#[must_use]
pub fn write_hex_bytes(bytes: &[u8]) -> String {
    let mut s = String::with_capacity(bytes.len() * 2);
    for byte in bytes {
        // Writing to a String is infallible
        let _ = write!(&mut s, "{byte:02x}");
    }
    s
}

// Calculate part size for multipart upload (max 5 GB)
// <https://docs.aws.amazon.com/AmazonS3/latest/userguide/qfacts.html>
/// # Errors
/// Will return `Err` if the file size exceeds 5 TB
/// or if the part size exceeds 5 GB
pub fn calculate_part_size(file_size: u64, buf_size: u64) -> Result<u64, S3Error> {
    log::info!("file size: {file_size}, buf size: {buf_size}");

    let mut part_size = buf_size.max(1);

    while {
        let calculated_parts = (file_size.saturating_add(part_size).saturating_sub(1)) / part_size;
        log::debug!(
            "Calculated parts: {calculated_parts}, Part size: {part_size}, Max parts: {MAX_PARTS_PER_UPLOAD}"
        );
        calculated_parts > MAX_PARTS_PER_UPLOAD as u64
    } {
        part_size = part_size.saturating_mul(2);
    }

    if part_size > MAX_PART_SIZE_BYTES {
        log::error!("max part size 5 GB");
        return Err(S3Error::Other("max part size 5 GB".to_string()));
    }

    log::info!("part size: {part_size}");

    Ok(part_size)
}

/// # Errors
///
/// Will return `Err` if can not open the file
pub fn blake3(file: &Path) -> std::io::Result<String> {
    let mut file = std::fs::File::open(file)?;
    let mut hasher = blake3::Hasher::new();
    let mut buf = vec![0_u8; 65_536];

    loop {
        let size = file.read(&mut buf)?;
        if size == 0 {
            break;
        }
        let chunk = buf
            .get(..size)
            .ok_or_else(|| Error::other(format!("invalid buffer slice length: {size}")))?;
        hasher.update(chunk);
    }

    Ok(hasher.finalize().to_hex().to_string())
}

/// Throttle downloads by introducing a delay based on bandwidth and chunk size.
///
/// # Arguments
/// - `bandwidth_kb`: Bandwidth in kilobytes per second (must be > 0).
/// - `chunk_size`: Size of the data chunk in bytes (must be > 0).
///
/// # Errors
/// - Returns `Err` if bandwidth is 0, chunk size is 0, bandwidth value overflows, or duration calculation is invalid.
///
/// # Returns
/// - `Ok(())` if throttling is successful.
pub async fn throttle_download(bandwidth_kb: usize, chunk_size: usize) -> std::io::Result<()> {
    if bandwidth_kb == 0 {
        return Err(Error::new(
            ErrorKind::InvalidInput,
            "Bandwidth cannot be 0 KB/s",
        ));
    }

    if chunk_size == 0 {
        return Err(Error::new(
            ErrorKind::InvalidInput,
            "Chunk size cannot be 0 bytes",
        ));
    }

    // Calculate bandwidth in bytes per second
    let bandwidth_bytes = bandwidth_kb
        .checked_mul(1024)
        .ok_or_else(|| Error::new(ErrorKind::InvalidInput, "Bandwidth value too large"))?;

    // Calculate duration in nanoseconds using integer arithmetic
    // duration = (chunk_size * 1_000_000_000) / bandwidth_bytes
    let nanos = u64::try_from(chunk_size)
        .ok()
        .and_then(|cs| cs.checked_mul(1_000_000_000))
        .and_then(|n| {
            u64::try_from(bandwidth_bytes)
                .ok()
                .and_then(|bb| n.checked_div(bb))
        })
        .ok_or_else(|| Error::new(ErrorKind::InvalidInput, "Duration calculation overflow"))?;

    let duration = Duration::from_nanos(nanos);

    // Introduce the delay for throttling
    sleep(duration).await;

    // Log the throttling details
    log::debug!(
        "Throttling download: Bandwidth = {bandwidth_kb} KB/s, Chunk size = {chunk_size} bytes, Sleep duration = {:.3} seconds",
        duration.as_secs_f64()
    );

    Ok(())
}

#[cfg(test)]
#[allow(
    clippy::unwrap_used,
    clippy::expect_used,
    clippy::panic,
    clippy::indexing_slicing,
    clippy::unnecessary_wraps,
    clippy::format_collect
)]
mod tests {
    use super::*;
    use crate::s3::limits::{MAX_OBJECT_SIZE_BYTES, MAX_PART_SIZE_BYTES, MAX_PARTS_PER_UPLOAD};
    use crate::stream::iterator::PartIterator;
    use std::io::Write;
    use tempfile::NamedTempFile;

    #[test]
    fn test_sha256_digest() {
        let digest = sha256_digest(b"hello world");
        assert_eq!(
            write_hex_bytes(digest.as_ref()),
            "b94d27b9934d3e08a52e52d7da7dabfac484efe37a5380ee9088f7ace2efcde9"
        );
    }

    #[test]
    fn test_base64_md5() {
        let md5 = base64_md5(b"hello world");
        assert_eq!(md5, "XrY7u+Ae7tCTyyK7j1rNww==");
    }

    #[test]
    fn test_sha256_hmac() {
        let key = b"key";
        let msg = b"The quick brown fox jumps over the lazy dog";
        let tag = sha256_hmac(key, msg);
        assert_eq!(
            write_hex_bytes(tag.as_ref()),
            "f7bc83f430538424b13298e6aa6fb143ef4d59a14946175997479dbc2d1a3cd8"
        );
    }

    #[test]
    fn test_write_hex_bytes() {
        let bytes = b"hello world";
        assert_eq!(write_hex_bytes(bytes), "68656c6c6f20776f726c64");
    }

    #[test]
    fn test_blake3() {
        let mut file = NamedTempFile::new().unwrap();
        file.write_all(b"hello world").unwrap();
        let hash = blake3(file.path()).unwrap();
        assert_eq!(
            hash,
            "d74981efa70a0c880b8d8c1985d075dbcbf679b99a5f9914e5aaf96b831a9e24"
        );
    }

    #[test]
    fn test_calculate_part_size() {
        let file_size = 5 * 1024 * 1024 * 1000 * 1000;
        let buf_size = 10 * 1024 * 1024;
        let part_size = calculate_part_size(file_size, buf_size).unwrap();
        assert!(part_size <= MAX_PART_SIZE_BYTES);
        assert_eq!(part_size, 671_088_640);
    }

    #[test]
    fn test_calculate_part_size_10000parts() {
        let file_size = 5 * 1024 * 1024 * 1000 * 1000;
        let buf_size = 512 * 1024 * 1024;
        let part_size = calculate_part_size(file_size, buf_size).unwrap();
        assert!(part_size <= MAX_PART_SIZE_BYTES);
        assert_eq!(part_size, buf_size);
    }

    #[test]
    fn test_calculate_part_size_within_limits() {
        let result = calculate_part_size(1_000_000, 1_000);
        assert!(result.is_ok());
    }

    #[test]
    fn test_calculate_part_size_exceeds_max_part_size() {
        let result = calculate_part_size(1_000_000, MAX_PART_SIZE_BYTES + 1);
        assert!(result.is_err());
        assert_eq!(result.err().unwrap().to_string(), "max part size 5 GB");
    }

    #[test]
    fn test_calculate_part_size_5tb_part_512mb() {
        let file_size = 5 * 1024 * 1024 * 1000 * 1000;
        let buf_size = 1;
        let result = calculate_part_size(file_size, buf_size).unwrap();
        assert_eq!(result, MAX_PART_SIZE_BYTES / 10);
    }

    #[test]
    fn test_calculate_part_size_5tb_part_1gb() {
        let file_size = 5 * 1024 * 1024 * 1000 * 1000;
        let buf_size = 1024 * 1024 * 1024;
        let result = calculate_part_size(file_size, buf_size).unwrap();
        assert_eq!(result, MAX_PART_SIZE_BYTES / 5);
    }

    #[test]
    fn test_calculate_part_size_max_part_size() {
        let file_size = 5 * 1024 * 1024 * 1000 * 1000;
        let part_size = calculate_part_size(file_size, MAX_PART_SIZE_BYTES).unwrap();
        assert_eq!(part_size, MAX_PART_SIZE_BYTES);
        let (number, seek, chunk) = PartIterator::new(file_size, part_size).last().unwrap();
        assert_eq!(file_size, seek + chunk);
        assert!(usize::from(number) <= MAX_PARTS_PER_UPLOAD);
    }

    #[test]
    fn test_calculate_part_size_max_part_size_512() {
        let file_size = 5 * 1024 * 1024 * 1000 * 1000;
        let buf_size = 512 * 1024 * 1024;
        let part_size = calculate_part_size(file_size, buf_size).unwrap();
        assert_eq!(part_size, buf_size);
        let (number, seek, chunk) = PartIterator::new(file_size, part_size).last().unwrap();
        assert_eq!(file_size, seek + chunk);
        assert!(usize::from(number) <= MAX_PARTS_PER_UPLOAD);
    }

    #[test]
    fn test_calculate_part_size_max_part_per_upload() {
        let buf_size = 549_755_814;
        let part_size = calculate_part_size(MAX_OBJECT_SIZE_BYTES, buf_size).unwrap();
        let (number, seek, chunk) = PartIterator::new(MAX_OBJECT_SIZE_BYTES, part_size)
            .last()
            .unwrap();
        assert_eq!(MAX_OBJECT_SIZE_BYTES, seek + chunk);
        assert!(usize::from(number) <= MAX_PARTS_PER_UPLOAD);
    }

    #[test]
    fn test_calculate_part_size_max_part_per_upload_1() {
        let buf_size = 1;
        let part_size = calculate_part_size(MAX_OBJECT_SIZE_BYTES, buf_size).unwrap();
        let (number, seek, chunk) = PartIterator::new(MAX_OBJECT_SIZE_BYTES, part_size)
            .last()
            .unwrap();
        assert_eq!(MAX_OBJECT_SIZE_BYTES, seek + chunk);
        assert!(usize::from(number) <= MAX_PARTS_PER_UPLOAD);
    }

    #[test]
    fn test_calculate_part_size_max_part_per_upload_2() {
        let part_size = calculate_part_size(15, 4).unwrap();
        let mut parts = PartIterator::new(15, part_size);
        assert_eq!(parts.next(), Some((1, 0, 4)));
        assert_eq!(parts.next(), Some((2, 4, 4)));
        assert_eq!(parts.next(), Some((3, 8, 4)));
        assert_eq!(parts.next(), Some((4, 12, 3)));
        assert_eq!(parts.next(), None);
    }

    #[test]
    fn test_calculate_part_size_max_part_per_upload_3() {
        let part_size = calculate_part_size(15, 5).unwrap();
        let mut parts = PartIterator::new(15, part_size);
        assert_eq!(parts.next(), Some((1, 0, 5)));
        assert_eq!(parts.next(), Some((2, 5, 5)));
        assert_eq!(parts.next(), Some((3, 10, 5)));
        assert_eq!(parts.next(), None);
    }

    #[test]
    fn test_calculate_part_size_max_part_per_upload_4() {
        let buf_size = 52_428_800;
        let part_size = calculate_part_size(524_288_000_000, buf_size).unwrap();
        let (number, seek, chunk) = PartIterator::new(524_288_000_000, part_size)
            .last()
            .unwrap();
        assert_eq!(524_288_000_000, seek + chunk);
        assert!(usize::from(number) <= MAX_PARTS_PER_UPLOAD);
    }

    #[test]
    fn test_calculate_part_size_10k_parts() {
        // Use a file size that results in exactly 10,000 parts
        let file_size = 100_000_000; // Assume a fixed file size
        let buf_size = 10_000; // Assume a fixed buffer size

        let part_size = calculate_part_size(file_size, buf_size).unwrap();

        // Ensure that the calculated part size leads to exactly 10,000 parts
        assert_eq!(file_size.div_ceil(part_size), 10_000);
    }

    #[test]
    fn test_calculate_part_size_less_than_10k_parts() {
        // Use a file size that results in less than 10,000 parts
        let file_size = 90_000_000; // Assume a fixed file size
        let buf_size = 10_000; // Set buf_size to a value that ensures less than 10,000 parts

        let part_size = calculate_part_size(file_size, buf_size).unwrap();

        // Ensure that the calculated part size results in less than 10,000 parts
        assert!(file_size.div_ceil(part_size) < 10_000);
    }

    #[test]
    fn test_calculate_part_size_buffer_size_0() {
        let part_size = calculate_part_size(MAX_OBJECT_SIZE_BYTES, 0).unwrap();
        let (number, seek, chunk) = PartIterator::new(MAX_OBJECT_SIZE_BYTES, part_size)
            .last()
            .unwrap();
        assert_eq!(MAX_OBJECT_SIZE_BYTES, seek + chunk);
        assert!(usize::from(number) <= MAX_PARTS_PER_UPLOAD);
        assert!((MAX_OBJECT_SIZE_BYTES.div_ceil(part_size)) < 10_000);
    }

    #[tokio::test]
    async fn test_throttle_download_valid() {
        let bandwidth_kb = 1; // 1 KB/s
        let chunk_size = 1024; // 1 KB

        // Test 1: Exact division (1024 bytes / 1024 bytes/sec = 1 sec)
        let start = std::time::Instant::now();
        let result = throttle_download(bandwidth_kb, chunk_size).await;

        assert!(result.is_ok());
        let elapsed = start.elapsed();
        assert!(elapsed.as_secs_f64() >= 1.0);

        // Test 2: Fractional seconds (512 bytes / 1024 bytes/sec = 0.5 sec)
        let start = std::time::Instant::now();
        let result = throttle_download(1, 512).await;

        assert!(result.is_ok());
        let elapsed = start.elapsed();
        assert!(elapsed.as_secs_f64() >= 0.5);
    }

    #[tokio::test]
    async fn test_throttle_download_zero_bandwidth() {
        let result = throttle_download(0, 1024).await;
        assert!(result.is_err());
        assert_eq!(
            result.unwrap_err().to_string(),
            "Bandwidth cannot be 0 KB/s"
        );
    }

    #[tokio::test]
    async fn test_throttle_download_zero_chunk_size() {
        let result = throttle_download(1, 0).await;
        assert!(result.is_err());
        assert_eq!(
            result.unwrap_err().to_string(),
            "Chunk size cannot be 0 bytes"
        );
    }

    #[tokio::test]
    async fn test_throttle_download_large_bandwidth() {
        let bandwidth_kb = usize::MAX / 1024 + 1; // Exceeds usize when multiplied by 1024
        let result = throttle_download(bandwidth_kb, 1024).await;
        assert!(result.is_err());
        assert_eq!(result.unwrap_err().to_string(), "Bandwidth value too large");
    }
}