// Copyright 2026 TiKV Project Authors. Licensed under Apache-2.0.

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

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

// TiDB transaction entries can be up to 120 MiB. Leave enough room for SST
// block encoding overhead while preventing corrupt files from advertising
// multi-gigabyte allocations.
pub(crate) const MAX_BLOCK_SIZE: usize = 128 * 1024 * 1024;
const MAX_ZSTD_WINDOW_LOG: u32 = 26;

/// Compression algorithms used by TiKV legacy Backup SST files.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[repr(u8)]
pub enum CompressionType {
    None = 0,
    Snappy = 1,
    Lz4 = 4,
    Zstd = 7,
}

impl TryFrom<u8> for CompressionType {
    type Error = Error;

    fn try_from(value: u8) -> Result<Self> {
        match value {
            0 => Ok(Self::None),
            1 => Ok(Self::Snappy),
            4 => Ok(Self::Lz4),
            7 => Ok(Self::Zstd),
            other => Err(Error::Unsupported(format!("compression type {other}"))),
        }
    }
}

impl CompressionType {
    pub(crate) fn rocksdb_name(self) -> &'static str {
        match self {
            Self::None => "NoCompression",
            Self::Snappy => "Snappy",
            Self::Lz4 => "LZ4",
            Self::Zstd => "ZSTD",
        }
    }
}

pub(crate) fn compress(
    ty: CompressionType,
    format_version: u32,
    input: &[u8],
    level: i32,
) -> Result<Vec<u8>> {
    match ty {
        CompressionType::None => Ok(input.to_vec()),
        CompressionType::Snappy => snap::raw::Encoder::new()
            .compress_vec(input)
            .map_err(|err| Error::Corruption(format!("snappy compression failed: {err}"))),
        CompressionType::Lz4 => {
            let mut output = Vec::new();
            if format_version >= 2 {
                coding::put_varint32(
                    &mut output,
                    u32::try_from(input.len())
                        .map_err(|_| Error::Unsupported("block larger than 4 GiB".into()))?,
                );
            } else {
                let len = u64::try_from(input.len())
                    .map_err(|_| Error::Unsupported("block length overflow".into()))?;
                output.extend_from_slice(&len.to_le_bytes());
            }
            output.extend_from_slice(&lz4::block::compress(input, None, false)?);
            Ok(output)
        }
        CompressionType::Zstd => {
            let mut output = Vec::new();
            coding::put_varint32(
                &mut output,
                u32::try_from(input.len())
                    .map_err(|_| Error::Unsupported("block larger than 4 GiB".into()))?,
            );
            output.extend_from_slice(&zstd::stream::encode_all(Cursor::new(input), level)?);
            Ok(output)
        }
    }
}

pub(crate) fn decompress(
    ty: CompressionType,
    format_version: u32,
    input: &[u8],
) -> Result<Vec<u8>> {
    match ty {
        CompressionType::None => Ok(input.to_vec()),
        CompressionType::Snappy => {
            let decompressed_len = snap::raw::decompress_len(input)
                .map_err(|err| Error::Corruption(format!("invalid snappy block: {err}")))?;
            validate_decompressed_len(decompressed_len)?;
            let mut output = vec![0; decompressed_len];
            let written = snap::raw::Decoder::new()
                .decompress(input, &mut output)
                .map_err(|err| Error::Corruption(format!("invalid snappy block: {err}")))?;
            output.truncate(written);
            Ok(output)
        }
        CompressionType::Lz4 => {
            let mut pos = 0;
            let decompressed_len = if format_version >= 2 {
                coding::read_varint32(input, &mut pos)?
            } else {
                let bytes = input
                    .get(..8)
                    .ok_or_else(|| Error::Corruption("truncated legacy LZ4 header".into()))?;
                pos = 8;
                u32::from_le_bytes(bytes[..4].try_into().unwrap())
            };
            validate_decompressed_len(decompressed_len as usize)?;
            lz4::block::decompress(
                input
                    .get(pos..)
                    .ok_or_else(|| Error::Corruption("truncated LZ4 block".into()))?,
                Some(decompressed_len as i32),
            )
            .map_err(Error::from)
        }
        CompressionType::Zstd => {
            let mut pos = 0;
            let expected_len = coding::read_varint32(input, &mut pos)? as usize;
            validate_decompressed_len(expected_len)?;
            let mut decoder = zstd::stream::read::Decoder::new(Cursor::new(&input[pos..]))?;
            decoder.window_log_max(MAX_ZSTD_WINDOW_LOG)?;
            let mut output = vec![0; expected_len];
            decoder.read_exact(&mut output).map_err(|err| {
                Error::Corruption(format!("invalid or truncated zstd block: {err}"))
            })?;
            let mut extra = [0; 1];
            if decoder.read(&mut extra)? != 0 {
                return Err(Error::Corruption(format!(
                    "zstd block exceeds advertised size {expected_len}"
                )));
            }
            Ok(output)
        }
    }
}

fn validate_decompressed_len(len: usize) -> Result<()> {
    if len > MAX_BLOCK_SIZE {
        return Err(Error::Corruption(format!(
            "decompressed block size {len} exceeds limit {MAX_BLOCK_SIZE}"
        )));
    }
    Ok(())
}

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

    #[test]
    fn rejects_oversized_decompressed_blocks() {
        let oversized = u32::try_from(MAX_BLOCK_SIZE + 1).unwrap();

        let mut snappy = Vec::new();
        coding::put_varint32(&mut snappy, oversized);
        assert!(matches!(
            decompress(CompressionType::Snappy, 5, &snappy),
            Err(Error::Corruption(message)) if message.contains("exceeds limit")
        ));

        let mut lz4 = Vec::new();
        coding::put_varint32(&mut lz4, oversized);
        assert!(matches!(
            decompress(CompressionType::Lz4, 5, &lz4),
            Err(Error::Corruption(message)) if message.contains("exceeds limit")
        ));

        let mut zstd = Vec::new();
        coding::put_varint32(&mut zstd, oversized);
        assert!(matches!(
            decompress(CompressionType::Zstd, 5, &zstd),
            Err(Error::Corruption(message)) if message.contains("exceeds limit")
        ));
    }
}