// Copyright 2026 PingCAP, Inc.

use std::convert::TryFrom;

use crate::table::{Error, Result};

pub(crate) fn lz4_compress_bound_checked(src_len: usize) -> Result<i32> {
    let src_len_i32 = i32::try_from(src_len).map_err(|_| Error::Lz4CompressBoundFailed {
        input_size: src_len,
    })?;
    let compress_bound = unsafe { lz4::liblz4::LZ4_compressBound(src_len_i32) };
    if compress_bound <= 0 {
        return Err(Error::Lz4CompressBoundFailed {
            input_size: src_len,
        });
    }
    Ok(compress_bound)
}

pub(crate) unsafe fn lz4_compress_checked(
    src: *const libc::c_char,
    src_len: i32,
    dst: *mut libc::c_char,
    dst_capacity: i32,
) -> Result<usize> {
    let size = lz4::liblz4::LZ4_compress_default(src, dst, src_len, dst_capacity);
    if size <= 0 {
        return Err(Error::Lz4CompressionFailed {
            input_size: src_len as usize,
        });
    }
    Ok(size as usize)
}

pub(crate) fn lz4_compress_append(src: &[u8], dst: &mut Vec<u8>) -> Result<()> {
    let src_len = src.len();
    let compress_bound = lz4_compress_bound_checked(src_len)?;
    // Safe because `lz4_compress_bound_checked` already verified `src_len <=
    // i32::MAX`.
    let src_len_i32 = src_len as i32;

    unsafe {
        let dst_off = dst.len();
        dst.reserve(compress_bound as usize);
        let dst_ptr = dst.as_mut_ptr().add(dst_off) as *mut libc::c_char;
        let size = lz4_compress_checked(
            src.as_ptr() as *const libc::c_char,
            src_len_i32,
            dst_ptr,
            compress_bound,
        )?;
        dst.set_len(dst_off + size);
    }
    Ok(())
}

pub(crate) fn zstd_compress_bound_checked(src_len: usize) -> Result<usize> {
    let compress_bound = unsafe { zstd_sys::ZSTD_compressBound(src_len) };
    if unsafe { zstd_sys::ZSTD_isError(compress_bound) } != 0u32 {
        let err_name_ptr = unsafe { zstd_sys::ZSTD_getErrorName(compress_bound) };
        let reason = unsafe { std::ffi::CStr::from_ptr(err_name_ptr) }
            .to_string_lossy()
            .into_owned();
        return Err(Error::ZstdCompressBoundFailed {
            input_size: src_len,
            reason,
        });
    }
    Ok(compress_bound)
}

pub(crate) unsafe fn zstd_compress_checked(
    src: *const libc::c_void,
    src_len: usize,
    dst: *mut libc::c_void,
    dst_capacity: usize,
    compression_lvl: libc::c_int,
) -> Result<usize> {
    let size = zstd_sys::ZSTD_compress(dst, dst_capacity, src, src_len, compression_lvl);
    if zstd_sys::ZSTD_isError(size) != 0u32 {
        let err_name_ptr = zstd_sys::ZSTD_getErrorName(size);
        let reason = std::ffi::CStr::from_ptr(err_name_ptr)
            .to_string_lossy()
            .into_owned();
        return Err(Error::ZstdCompressionFailed {
            input_size: src_len,
            reason,
        });
    }
    Ok(size)
}

pub(crate) fn zstd_compress_append(
    src: &[u8],
    compression_lvl: i32,
    dst: &mut Vec<u8>,
) -> Result<()> {
    let src_len = src.len();
    let compress_bound = zstd_compress_bound_checked(src_len)?;
    unsafe {
        let dst_off = dst.len();
        dst.reserve(compress_bound);
        let dst_ptr = dst.as_mut_ptr().add(dst_off) as *mut libc::c_void;
        let size = zstd_compress_checked(
            src.as_ptr() as *const libc::c_void,
            src_len,
            dst_ptr,
            compress_bound,
            compression_lvl as libc::c_int,
        )?;
        dst.set_len(dst_off + size);
    }
    Ok(())
}

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

    #[test]
    fn test_zstd_compress_error_is_handled() {
        let src = vec![0u8; 1000];
        let mut dst = Vec::<u8>::with_capacity(1);
        let err = unsafe {
            zstd_compress_checked(
                src.as_ptr() as *const libc::c_void,
                src.len(),
                dst.as_mut_ptr() as *mut libc::c_void,
                dst.capacity(),
                3,
            )
            .unwrap_err()
        };
        match err {
            Error::ZstdCompressionFailed { input_size, .. } => {
                assert_eq!(input_size, src.len());
            }
            other => panic!("unexpected error: {:?}", other),
        }
    }

    #[test]
    fn test_zstd_compress_bound_error_is_handled() {
        let max_input_size = if usize::BITS == 64 {
            0xFF00FF00FF00FF00usize
        } else {
            0xFF00FF00usize
        };
        let err = zstd_compress_bound_checked(max_input_size).unwrap_err();
        match err {
            Error::ZstdCompressBoundFailed { input_size, .. } => {
                assert_eq!(input_size, max_input_size);
            }
            other => panic!("unexpected error: {:?}", other),
        }
    }

    #[test]
    fn test_zstd_compress_append_appends_output() {
        let src = vec![0u8; 1024];
        let prefix = vec![7u8, 8u8, 9u8];
        let mut dst = prefix.clone();

        zstd_compress_append(&src, 3, &mut dst).unwrap();

        assert!(dst.len() > prefix.len());
        assert_eq!(&dst[..prefix.len()], prefix.as_slice());
        let compressed_len = dst.len() - prefix.len();
        assert!(compressed_len <= zstd_compress_bound_checked(src.len()).unwrap());
    }

    #[test]
    fn test_lz4_compress_error_is_handled() {
        let src = vec![0u8; 1000];
        let mut dst = Vec::<u8>::with_capacity(1);
        let err = unsafe {
            lz4_compress_checked(
                src.as_ptr() as *const libc::c_char,
                src.len() as i32,
                dst.as_mut_ptr() as *mut libc::c_char,
                dst.capacity() as i32,
            )
            .unwrap_err()
        };
        match err {
            Error::Lz4CompressionFailed { input_size } => {
                assert_eq!(input_size, src.len());
            }
            other => panic!("unexpected error: {:?}", other),
        }
    }

    #[test]
    fn test_lz4_compress_append_appends_output() {
        let src = vec![0u8; 1024];
        let prefix = vec![1u8, 2u8, 3u8];
        let mut dst = prefix.clone();

        lz4_compress_append(&src, &mut dst).unwrap();

        assert!(dst.len() > prefix.len());
        assert_eq!(&dst[..prefix.len()], prefix.as_slice());
        let compressed_len = dst.len() - prefix.len();
        assert!((compressed_len as i32) <= lz4_compress_bound_checked(src.len()).unwrap());
    }

    #[test]
    fn test_lz4_compress_bound_error_is_handled() {
        let too_large = i32::MAX as usize + 1;
        let err = lz4_compress_bound_checked(too_large).unwrap_err();
        match err {
            Error::Lz4CompressBoundFailed { input_size } => {
                assert_eq!(input_size, too_large);
            }
            other => panic!("unexpected error: {:?}", other),
        }
    }
}