datum-mq 0.10.9

use std::io::{Cursor, Read, Write};

use lz4_flex::frame::{BlockMode, BlockSize, FrameDecoder, FrameEncoder, FrameInfo};
use zstd::bulk::Compressor as ZstdCompressor;
use zstd::zstd_safe::CParameter;

use crate::native::{KafkaClientError, KafkaClientResult};

/// Hard ceiling for one expanded Kafka record batch.
///
/// Kafka's default Fetch response ceiling is 50 MiB. Keeping the per-batch
/// limit slightly above it supports that default while bounding malicious
/// frame expansion and the owned buffer retained during record parsing.
pub(crate) const MAX_DECOMPRESSED_RECORD_BATCH_BYTES: usize = 64 * 1024 * 1024;

const LZ4_ATTRIBUTE: i16 = 3;
const ZSTD_ATTRIBUTE: i16 = 4;
const ZSTD_WINDOW_LOG_MAX: u32 = 26;

#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub(crate) enum CompressionCodec {
    #[default]
    None,
    Lz4,
    Zstd,
}

impl CompressionCodec {
    pub(crate) fn from_attribute(attribute: i16) -> KafkaClientResult<Self> {
        match attribute {
            0 => Ok(Self::None),
            LZ4_ATTRIBUTE => Ok(Self::Lz4),
            ZSTD_ATTRIBUTE => Ok(Self::Zstd),
            1 => Err(KafkaClientError::unsupported(
                "gzip-compressed Kafka record batches are not supported",
            )),
            2 => Err(KafkaClientError::unsupported(
                "Snappy-compressed Kafka record batches are not supported",
            )),
            other => Err(KafkaClientError::unsupported(format!(
                "unknown Kafka record-batch compression codec {other}"
            ))),
        }
    }

    pub(crate) fn attribute(self) -> i16 {
        match self {
            Self::None => 0,
            Self::Lz4 => LZ4_ATTRIBUTE,
            Self::Zstd => ZSTD_ATTRIBUTE,
        }
    }

    pub(crate) fn name(self) -> &'static str {
        match self {
            Self::None => "none",
            Self::Lz4 => "lz4",
            Self::Zstd => "zstd",
        }
    }
}

/// A reusable Zstd compression context for one producer's produce loop.
///
/// `zstd::bulk::Compressor` wraps a `CCtx` whose internal match-finder tables
/// are the expensive-to-allocate part of Zstd compression; reusing it across
/// batches (instead of building a fresh streaming `Encoder` per batch) avoids
/// re-allocating that state every call. The produce loop is single-owner, so
/// holding this by value (no lock, no thread-local) is sufficient.
pub(crate) struct ZstdEncoderContext {
    compressor: ZstdCompressor<'static>,
}

impl ZstdEncoderContext {
    pub(crate) fn new() -> KafkaClientResult<Self> {
        let mut compressor = ZstdCompressor::new(zstd::DEFAULT_COMPRESSION_LEVEL)
            .map_err(|error| compression_error(CompressionCodec::Zstd, error))?;
        // Sticky parameters on the CCtx: `ZSTD_compress2` (what `Compressor::compress`
        // calls) always starts a fresh frame per call but keeps these set across calls.
        compressor
            .set_parameter(CParameter::ChecksumFlag(false))
            .map_err(|error| compression_error(CompressionCodec::Zstd, error))?;
        compressor
            .set_parameter(CParameter::ContentSizeFlag(false))
            .map_err(|error| compression_error(CompressionCodec::Zstd, error))?;
        Ok(Self { compressor })
    }
}

pub(crate) fn compress_records(
    codec: CompressionCodec,
    records: Vec<u8>,
    zstd_context: &mut ZstdEncoderContext,
) -> KafkaClientResult<Vec<u8>> {
    match codec {
        CompressionCodec::None => Ok(records),
        CompressionCodec::Lz4 => compress_lz4(&records),
        CompressionCodec::Zstd => compress_zstd(&records, zstd_context),
    }
}

pub(crate) fn decompress_records(
    codec: CompressionCodec,
    records: &[u8],
) -> KafkaClientResult<Vec<u8>> {
    decompress_records_with_limit(codec, records, MAX_DECOMPRESSED_RECORD_BATCH_BYTES)
}

fn compress_lz4(records: &[u8]) -> KafkaClientResult<Vec<u8>> {
    // Kafka's Java client writes standard LZ4 frames with independent 64 KiB
    // blocks and no block/content checksum or content-size field.
    let frame = FrameInfo::new()
        .block_size(BlockSize::Max64KB)
        .block_mode(BlockMode::Independent);
    let mut encoder = FrameEncoder::with_frame_info(frame, Vec::new());
    encoder
        .write_all(records)
        .map_err(|error| compression_error(CompressionCodec::Lz4, error))?;
    encoder
        .finish()
        .map_err(|error| compression_error(CompressionCodec::Lz4, error))
}

fn compress_zstd(records: &[u8], context: &mut ZstdEncoderContext) -> KafkaClientResult<Vec<u8>> {
    context
        .compressor
        .compress(records)
        .map_err(|error| compression_error(CompressionCodec::Zstd, error))
}

fn decompress_records_with_limit(
    codec: CompressionCodec,
    records: &[u8],
    limit: usize,
) -> KafkaClientResult<Vec<u8>> {
    match codec {
        CompressionCodec::None => Ok(records.to_vec()),
        CompressionCodec::Lz4 => read_bounded(
            codec,
            FrameDecoder::new(Cursor::new(records)),
            records.len(),
            limit,
        ),
        CompressionCodec::Zstd => {
            let declared_size = zstd::zstd_safe::get_frame_content_size(records)
                .map_err(|error| compression_error(codec, error))?;
            if declared_size.is_some_and(|size| size > limit as u64) {
                return Err(KafkaClientError::decompression_limit(codec.name(), limit));
            }
            let mut decoder = zstd::stream::read::Decoder::new(Cursor::new(records))
                .map_err(|error| compression_error(codec, error))?
                .single_frame();
            decoder
                .window_log_max(ZSTD_WINDOW_LOG_MAX)
                .map_err(|error| compression_error(codec, error))?;
            let output = read_bounded(codec, &mut decoder, records.len(), limit)?;
            let reader = decoder.finish();
            let consumed = reader
                .get_ref()
                .position()
                .saturating_sub(reader.buffer().len() as u64);
            if consumed != records.len() as u64 {
                return Err(KafkaClientError::decompression(
                    codec.name(),
                    "trailing bytes after the Zstd frame",
                ));
            }
            Ok(output)
        }
    }
}

fn read_bounded(
    codec: CompressionCodec,
    mut reader: impl Read,
    compressed_len: usize,
    limit: usize,
) -> KafkaClientResult<Vec<u8>> {
    let initial_capacity = compressed_len.saturating_mul(4).min(limit);
    let mut output = Vec::with_capacity(initial_capacity);
    let mut chunk = [0_u8; 16 * 1024];
    loop {
        let read = reader
            .read(&mut chunk)
            .map_err(|error| compression_error(codec, error))?;
        if read == 0 {
            return Ok(output);
        }
        if output.len().saturating_add(read) > limit {
            return Err(KafkaClientError::decompression_limit(codec.name(), limit));
        }
        output.extend_from_slice(&chunk[..read]);
    }
}

fn compression_error(codec: CompressionCodec, error: impl std::fmt::Display) -> KafkaClientError {
    KafkaClientError::decompression(codec.name(), error.to_string())
}

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

    fn compress(codec: CompressionCodec, records: Vec<u8>) -> KafkaClientResult<Vec<u8>> {
        let mut zstd_context = ZstdEncoderContext::new().expect("zstd encoder context");
        compress_records(codec, records, &mut zstd_context)
    }

    #[test]
    fn lz4_and_zstd_round_trip_standard_frames() {
        let records = b"record-data-record-data-record-data";
        for codec in [CompressionCodec::Lz4, CompressionCodec::Zstd] {
            let encoded = compress(codec, records.to_vec()).expect("compress records");
            let decoded = decompress_records(codec, &encoded).expect("decompress records");
            assert_eq!(decoded, records, "codec={}", codec.name());
        }
    }

    #[test]
    fn zstd_context_reuse_produces_independent_valid_frames() {
        let mut zstd_context = ZstdEncoderContext::new().expect("zstd encoder context");
        let batches: [&[u8]; 3] = [b"first-batch", b"a-different-second-batch", b"third"];
        for batch in batches {
            let encoded =
                compress_records(CompressionCodec::Zstd, batch.to_vec(), &mut zstd_context)
                    .expect("compress with reused context");
            let decoded =
                decompress_records(CompressionCodec::Zstd, &encoded).expect("decompress records");
            assert_eq!(decoded, batch);
        }
    }

    #[test]
    fn decompression_expansion_limit_is_typed() {
        let records = vec![b'x'; 4_096];
        for codec in [CompressionCodec::Lz4, CompressionCodec::Zstd] {
            let encoded = compress(codec, records.clone()).expect("compress records");
            let error = decompress_records_with_limit(codec, &encoded, 1_024)
                .expect_err("expansion limit must fail");
            assert!(matches!(
                error,
                KafkaClientError::DecompressionLimitExceeded {
                    codec: _,
                    limit: 1_024
                }
            ));
        }
    }

    #[test]
    fn malformed_frames_are_typed_and_never_panic() {
        for codec in [CompressionCodec::Lz4, CompressionCodec::Zstd] {
            let encoded = compress(codec, b"record-data".to_vec()).expect("compress records");
            let malformed = [&encoded[..encoded.len() / 2], &[0xa5; 16][..]];
            for bytes in malformed {
                let error = decompress_records_with_limit(codec, bytes, 1_024)
                    .expect_err("malformed frame must fail");
                assert!(matches!(error, KafkaClientError::Decompression { .. }));
            }
        }

        let mut zstd_with_garbage = compress(CompressionCodec::Zstd, b"record-data".to_vec())
            .expect("compress zstd records");
        zstd_with_garbage.extend_from_slice(b"garbage");
        let error =
            decompress_records_with_limit(CompressionCodec::Zstd, &zstd_with_garbage, 1_024)
                .expect_err("trailing garbage must fail");
        assert!(matches!(error, KafkaClientError::Decompression { .. }));
    }

    #[test]
    fn wrong_lz4_content_size_is_typed() {
        let frame = FrameInfo::new()
            .block_size(BlockSize::Max64KB)
            .block_mode(BlockMode::Independent)
            .content_size(Some(128));
        let mut encoder = FrameEncoder::with_frame_info(frame, Vec::new());
        encoder.write_all(b"short").expect("write lz4 input");
        encoder
            .try_finish()
            .expect_err("encoder must reject the wrong declared size");
        let mut encoded = encoder.get_ref().clone();
        encoded.extend_from_slice(&[0_u8; 4]);
        let error = decompress_records_with_limit(CompressionCodec::Lz4, &encoded, 1_024)
            .expect_err("decoder must reject the wrong declared size");
        assert!(matches!(error, KafkaClientError::Decompression { .. }));
    }
}