// Copyright 2021 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.

use std::{cmp::Ordering, convert::TryFrom, mem, ops::Deref};

use byteorder::{ByteOrder, LittleEndian};
use bytes::{Buf, BufMut};
use cloud_encryption::EncryptionKey;
use farmhash;
use txn_types::decode_lock_max_ts;
use xorf::BinaryFuse8;

use super::super::table::Value;
use crate::table::{
    BIT_HAS_OLD_VERSION, ChecksumType, Error, InnerKey, LZ4_COMPRESSION, NO_COMPRESSION, Result,
    SnapVersion, VALUE_VERSION_LEN, ZSTD_COMPRESSION,
    blobtable::BlobRef,
    compression::{lz4_compress_append, zstd_compress_append},
};

pub const PROP_KEY_SMALLEST: &[u8] = b"smallest";
pub const PROP_KEY_BIGGEST: &[u8] = b"biggest";
pub const PROP_KEY_MAX_TS: &[u8] = b"max_ts";
pub const PROP_KEY_MAX_LOCK_TS: &[u8] = b"max_lock_ts";
pub const PROP_KEY_ENTRIES: &[u8] = b"entries";
pub const PROP_KEY_OLD_ENTRIES: &[u8] = b"old_entries";
pub const PROP_KEY_TOMBS: &[u8] = b"tombs";
pub const PROP_KEY_KV_SIZE: &[u8] = b"kv_size";
pub const PROP_KEY_IN_USE_TOTAL_BLOB_SIZE: &[u8] = b"in_use_total_blob_size";
pub const PROP_KEY_ENCRYPTION_VER: &[u8] = b"encryption_ver";
pub const PROP_KEY_SNAP_VERSION: &[u8] = b"l0_ver"; // keep the l0_ver for compatibility.
pub const AUX_INDEX_BINARY_FUSE8: u32 = 1;
pub const INDEX_FORMAT_V1: u32 = 1;
pub const BLOCK_FORMAT_V1: u32 = 1;
pub const TABLE_FORMAT_V1: u16 = 1;
pub const MAGIC_NUMBER: u32 = 2940551257;
pub const MAGIC_NUMBER_SPLIT_L0: u32 = 2940551258;
pub const BLOCK_ADDR_SIZE: usize = mem::size_of::<BlockAddress>();
/// Safety limit for the **uncompressed** size of an old-version block.
///
/// Old versions of the same key are appended into the same old block and the
/// block is only rotated when the key changes. Under workloads where one key
/// has a very large version history, the uncompressed old block can grow very
/// large and exceed LZ4's supported input size, causing block corruption.
/// See https://github.com/tidbcloud/cloud-storage-engine/issues/4615.
pub const DEFAULT_OLD_BLOCK_UNCOMPRESSED_SIZE_LIMIT: usize = 1536 * 1024 * 1024; // 1.5 GiB

#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
pub struct TableBuilderOptions {
    pub block_size: usize,
    pub max_table_size: usize,
    pub compression_tps: [u8; 3],
    pub compression_lvl: i32,
    pub flush_split_l0: bool,
}

impl Default for TableBuilderOptions {
    fn default() -> Self {
        Self {
            block_size: 64 * 1024,
            max_table_size: 16 * 1024 * 1024,
            compression_tps: [LZ4_COMPRESSION, ZSTD_COMPRESSION, ZSTD_COMPRESSION],
            compression_lvl: 3,
            flush_split_l0: false,
        }
    }
}

/// A structure to manage a slice of entries in a block.
///
/// This struct is used to store the serialized data of entries and their
/// corresponding offsets. It provides methods to append new entries,
/// retrieve specific entries, and manage the internal buffer size.
#[derive(Default)]
pub(crate) struct EntrySlice {
    /// The buffer that holds the serialized entry data.
    buf: Vec<u8>,

    /// A vector that stores the end offsets of each entry in the buffer.
    /// This allows for efficient retrieval of individual entries.
    end_offs: Vec<u32>,
}

#[allow(dead_code)]
impl EntrySlice {
    pub(crate) fn append(&mut self, data: &[u8]) {
        self.buf.extend_from_slice(data);
        self.end_offs.push(self.buf.len() as u32);
    }

    fn append_value(&mut self, val: Value, blob_ref: Option<BlobRef>) {
        let old_len = self.buf.len();
        let new_len = if blob_ref.is_some() {
            old_len + val.encoded_size_with_blob_ref()
        } else {
            old_len + val.encoded_size()
        };
        self.buf.resize(new_len, 0);
        let slice = self.buf.as_mut_slice();
        match blob_ref {
            Some(blob_ref) => val.encode_with_blob_ref(&mut slice[old_len..], blob_ref),
            None => val.encode(&mut slice[old_len..]),
        }
        self.end_offs.push(new_len as u32);
    }

    pub(crate) fn length(&self) -> usize {
        self.end_offs.len()
    }

    pub(crate) fn get_last(&self) -> &[u8] {
        self.get_entry(self.length() - 1)
    }

    pub(crate) fn get_entry(&self, i: usize) -> &[u8] {
        let start_off = if i > 0 {
            self.end_offs[i - 1] as usize
        } else {
            0
        };
        let slice = self.buf.as_slice();
        &slice[start_off..self.end_offs[i] as usize]
    }

    pub(crate) fn size(&self) -> usize {
        self.buf.len() + self.end_offs.len() * 4
    }

    pub(crate) fn reset(&mut self) {
        self.buf.truncate(0);
        self.end_offs.truncate(0);
    }
}

/// A struct that manages the construction of a sstable format.
///
/// For the sstable format specification, refer to the comment above `SsTable`.
pub struct Builder {
    /// The unique identifier for the SSTable being built.
    sst_fid: u64,
    /// Responsible for constructing the current data block, handling the latest
    /// versions of keys.
    block_builder: BlockBuilder,
    /// Responsible for constructing the old data block, handling older versions
    /// of keys.
    old_builder: BlockBuilder,
    /// The target size for each block within the SSTable.
    block_size: usize,
    /// The type of checksum used for data integrity verification.
    checksum_type: ChecksumType,
    /// A collection of hash values for the keys, used for building auxiliary
    /// structures like Bloom filters.
    key_hashes: Vec<u64>,
    /// The smallest key in the SSTable, used for indexing and range queries.
    smallest: Vec<u8>,
    /// The largest key in the SSTable, used for indexing and range queries.
    biggest: Vec<u8>,
    /// The maximum timestamp found in the entries, used for versioning and
    /// time-travel queries.
    max_ts: u64,
    /// The count of old entries stored in the old data block.
    old_entries: u32,
    /// The count of tombstone entries, representing deleted keys.
    tombs: u32,
    /// The total size of key-value pairs before compression, excluding metadata
    /// and version fields.
    kv_size: u64,
    /// The total size of values stored in the blob table that are currently in
    /// use.
    total_blob_size: u64,
    /// Optional encryption key used for encrypting the data blocks.
    encryption_key: Option<EncryptionKey>,
    snap_version: SnapVersion,
    is_lock_cf: bool,
    max_lock_ts: u64,

    /// The maximum allowed **uncompressed** size for the current old-version
    /// block.
    ///
    /// This limit exists to avoid generating blocks that exceed compression
    /// library input limits (e.g., LZ4) and to prevent producing malformed
    /// blocks on disk.
    old_block_uncompressed_size_limit: usize,
}

impl Default for Builder {
    fn default() -> Self {
        Self {
            sst_fid: 0,
            block_builder: BlockBuilder::default(),
            old_builder: BlockBuilder::default(),
            block_size: 0,
            checksum_type: ChecksumType::default(),
            key_hashes: Vec::new(),
            smallest: Vec::new(),
            biggest: Vec::new(),
            max_ts: 0,
            old_entries: 0,
            tombs: 0,
            kv_size: 0,
            total_blob_size: 0,
            encryption_key: None,
            snap_version: SnapVersion::default(),
            is_lock_cf: false,
            max_lock_ts: 0,
            old_block_uncompressed_size_limit: DEFAULT_OLD_BLOCK_UNCOMPRESSED_SIZE_LIMIT,
        }
    }
}

impl Builder {
    // compression_lvl is the compression level for zstd compression only.
    pub fn new(
        sst_fid: u64,
        block_size: usize,
        compression_tp: u8,
        compression_lvl: i32,
        checksum_type: ChecksumType,
        encryption_key: Option<EncryptionKey>,
    ) -> Self {
        let mut x = Self::default();
        x.sst_fid = sst_fid;
        x.checksum_type = checksum_type;
        x.block_size = block_size;
        x.block_builder.compression_tp = compression_tp;
        x.block_builder.compression_lvl = compression_lvl;
        x.old_builder.compression_tp = compression_tp;
        x.old_builder.compression_lvl = compression_lvl;
        x.encryption_key = encryption_key;
        x
    }

    pub fn set_snap_version(&mut self, snap_version: SnapVersion) {
        self.snap_version = snap_version;
    }

    pub fn set_is_lock_cf(&mut self) {
        self.is_lock_cf = true;
    }

    pub fn reset(&mut self, sst_fid: u64) {
        self.sst_fid = sst_fid;
        self.block_builder.reset_all();
        self.old_builder.reset_all();
        self.key_hashes.truncate(0);
        self.smallest.truncate(0);
        self.biggest.truncate(0);
        self.max_ts = 0;
        self.max_lock_ts = 0;
        self.tombs = 0;
        self.total_blob_size = 0;
        self.old_entries = 0;
        self.kv_size = 0;
    }

    fn add_property(buf: &mut Vec<u8>, key: &[u8], val: &[u8]) {
        buf.put_u16_le(key.len() as u16);
        buf.put_slice(key);
        buf.put_u32_le(val.len() as u32);
        buf.put_slice(val);
    }

    // Keys must be added in increasing order; otherwise, an out-of-order error will
    // be returned.
    pub fn add(
        &mut self,
        inner_key: InnerKey<'_>,
        val: &Value,
        blob_ref: Option<BlobRef>,
    ) -> Result<()> {
        let key = inner_key.deref();
        match self.block_builder.last_key_order(key) {
            Some(Ordering::Equal) => {
                self.block_builder
                    .set_last_entry_old_ver_if_zero(val.version);
                self.old_builder.add_entry(key, *val, blob_ref);
                let uncompressed_size = self.old_builder.uncompressed_block_size();
                if uncompressed_size > self.old_block_uncompressed_size_limit {
                    return Err(Error::OldVersionBlockTooLarge {
                        size: uncompressed_size,
                        limit: self.old_block_uncompressed_size_limit,
                    });
                }
                if let Some(blob_ref) = blob_ref {
                    self.total_blob_size += blob_ref.len as u64;
                } else if val.is_blob_ref() {
                    self.total_blob_size += val.get_blob_ref().len as u64;
                }
                self.old_entries += 1;
            }
            Some(Ordering::Greater) => {
                return Err(Error::OutOfOrder(
                    self.block_builder.block.tmp_keys.get_last().to_vec(),
                    key.to_vec(),
                ));
            }
            _ => {
                // Only try to finish block when the key is different than last.
                if self.block_builder.need_finish_block(self.block_size) {
                    self.block_builder
                        .finish_block(self.sst_fid, self.checksum_type)?;
                }
                if self.old_builder.need_finish_block(self.block_size) {
                    self.old_builder
                        .finish_block(self.sst_fid, self.checksum_type)?;
                }
                self.kv_size += (key.len() + val.user_meta_len()) as u64;
                if let Some(blob_ref) = blob_ref {
                    self.total_blob_size += blob_ref.len as u64;
                    self.kv_size += blob_ref.original_len as u64;
                } else if val.is_blob_ref() {
                    self.total_blob_size += val.get_blob_ref().len as u64;
                    self.kv_size += val.get_blob_ref().original_len as u64;
                } else {
                    self.kv_size += val.value_len() as u64;
                }
                self.block_builder.add_entry(key, *val, blob_ref);
                self.key_hashes.push(farmhash::fingerprint64(key));
                if self.smallest.is_empty() {
                    self.smallest.extend_from_slice(key);
                }
                if self.max_ts < val.version {
                    self.max_ts = val.version;
                }
                // Only count tombstones of the latest (main-region) version of a
                // key. Old-version entries go into `old_builder` and are tracked
                // separately by `old_entries`; counting their tombstones here
                // would make `tombs` inconsistent with `entries` (which only
                // covers the main region) and skew the tombstone-ratio used to
                // trigger major compaction. See issue #5409.
                if val.value_len() == 0 {
                    self.tombs += 1;
                }
            }
        }
        if self.is_lock_cf && !val.get_value().is_empty() {
            // Lock-CF files in generic engine tests may contain arbitrary bytes
            // rather than MVCC lock encodings. Preserve the historical
            // parse_ts-based behavior here: ignore malformed lock payloads when
            // collecting the optional max_lock_ts SST property.
            if let Ok(Some(ts)) = decode_lock_max_ts(val.get_value()) {
                let ts = ts.into_inner();
                if self.max_lock_ts < ts {
                    self.max_lock_ts = ts;
                }
            }
        }
        Ok(())
    }

    pub fn estimated_size(&self) -> usize {
        let mut size = self.block_builder.buf.len()
            + self.old_builder.buf.len()
            + self.block_builder.block_size()
            + self.old_builder.block_size();
        size += size / 32; // reserve extra capacity to avoid reallocate.
        size
    }

    fn encrypt_blocks(
        base_off: u32,
        encryption_key: &EncryptionKey,
        block_builder: &BlockBuilder,
        data_buf: &mut Vec<u8>,
    ) {
        let timer = tikv_util::time::Instant::now_coarse();
        for (i, addr) in block_builder.block_addrs.iter().enumerate() {
            let start = addr.curr_off as usize;
            let end = if i + 1 < block_builder.block_addrs.len() {
                block_builder.block_addrs[i + 1].curr_off as usize
            } else {
                block_builder.buf.len()
            };
            let data = &block_builder.buf[start..end];
            encryption_key.encrypt(data, addr.origin_fid, addr.curr_off + base_off, data_buf);
        }
        if !encryption_key.is_noop() {
            crate::metrics::ENGINE_SST_CRYPT_DURATION_STATIC
                .encrypt_data_blocks
                .observe(timer.saturating_elapsed().as_secs_f64());
        }
    }

    pub fn finish(&mut self, base_off: u32, data_buf: &mut Vec<u8>) -> Result<BuildResult> {
        if self.block_builder.block.kv_size > 0 {
            let last_key = self.block_builder.block.tmp_keys.get_last();
            self.biggest.extend_from_slice(last_key);
            self.block_builder
                .finish_block(self.sst_fid, self.checksum_type)?;
        }
        if self.old_builder.block.kv_size > 0 {
            self.old_builder
                .finish_block(self.sst_fid, self.checksum_type)?;
        }
        assert_eq!(self.block_builder.block_keys.length() > 0, true);
        if let Some(encryption_key) = &self.encryption_key {
            Self::encrypt_blocks(base_off, encryption_key, &self.block_builder, data_buf);
        } else {
            data_buf.extend_from_slice(self.block_builder.buf.as_slice());
        }
        let data_section_size = self.block_builder.buf.len() as u32;
        if let Some(encryption_key) = &self.encryption_key {
            Self::encrypt_blocks(
                base_off + data_section_size,
                encryption_key,
                &self.old_builder,
                data_buf,
            );
        } else {
            data_buf.extend_from_slice(self.old_builder.buf.as_slice());
        }
        let old_data_section_size = self.old_builder.buf.len() as u32;

        self.block_builder.build_index(base_off, self.checksum_type);
        data_buf.extend_from_slice(self.block_builder.buf.as_slice());
        let index_section_size = self.block_builder.buf.len() as u32;
        self.old_builder
            .build_index(base_off + data_section_size, self.checksum_type);
        data_buf.extend_from_slice(self.old_builder.buf.as_slice());
        let old_index_section_size = self.old_builder.buf.len() as u32;
        let aux_index_section_size = if let Ok(filter) = BinaryFuse8::try_from(&self.key_hashes) {
            let bin = filter.to_vec();
            let origin_len = data_buf.len();
            self.build_aux_index(data_buf, &bin);
            (data_buf.len() - origin_len) as u32
        } else {
            warn!("failed to build binary fuse 8 filter");
            0
        };
        self.build_properties(data_buf);

        let mut footer = Footer::default();
        footer.old_data_offset = data_section_size;
        footer.index_offset = footer.old_data_offset + old_data_section_size;
        footer.old_index_offset = footer.index_offset + index_section_size;
        footer.aux_index_offset = footer.old_index_offset + old_index_section_size;
        footer.properties_offset = footer.aux_index_offset + aux_index_section_size;
        footer.compression_type = self.block_builder.compression_tp;
        footer.checksum_type = self.checksum_type.value();
        footer.table_format_version = TABLE_FORMAT_V1;
        footer.magic = if self.snap_version.is_not_zero() {
            MAGIC_NUMBER_SPLIT_L0
        } else {
            MAGIC_NUMBER
        };
        footer.marshal(data_buf);

        Ok(BuildResult {
            id: self.sst_fid,
            meta_offset: footer.meta_offset(),
            smallest: self.smallest.clone(),
            biggest: self.biggest.clone(),
            #[cfg(test)]
            tiny_meta_offset: footer.tiny_meta_offset(),
        })
    }

    fn build_aux_index(&self, buf: &mut Vec<u8>, fuse8: &[u8]) {
        let origin_len = buf.len();
        buf.put_u32_le(0);
        buf.put_u32_le(AUX_INDEX_BINARY_FUSE8);
        buf.put_u32_le(fuse8.len() as u32);
        buf.extend_from_slice(fuse8);
        let checksum = self.checksum_type.checksum(&buf[(origin_len + 4)..]);
        LittleEndian::write_u32(&mut buf[origin_len..], checksum);
    }

    fn build_properties(&self, buf: &mut Vec<u8>) {
        let origin_len = buf.len();
        buf.put_u32_le(0);
        Builder::add_property(buf, PROP_KEY_SMALLEST, self.smallest.as_slice());
        Builder::add_property(buf, PROP_KEY_BIGGEST, self.biggest.as_slice());
        Builder::add_property(buf, PROP_KEY_MAX_TS, &self.max_ts.to_le_bytes());
        if self.max_lock_ts > 0 {
            Builder::add_property(buf, PROP_KEY_MAX_LOCK_TS, &self.max_lock_ts.to_le_bytes());
        }
        let entries = self.key_hashes.len() as u32;
        Builder::add_property(buf, PROP_KEY_ENTRIES, &entries.to_le_bytes());
        Builder::add_property(buf, PROP_KEY_OLD_ENTRIES, &self.old_entries.to_le_bytes());
        Builder::add_property(buf, PROP_KEY_TOMBS, &self.tombs.to_le_bytes());
        Builder::add_property(buf, PROP_KEY_KV_SIZE, &self.kv_size.to_le_bytes());
        Builder::add_property(
            buf,
            PROP_KEY_IN_USE_TOTAL_BLOB_SIZE,
            &self.total_blob_size.to_le_bytes(),
        );
        if let Some(encryption_key) = &self.encryption_key {
            Builder::add_property(
                buf,
                PROP_KEY_ENCRYPTION_VER,
                &encryption_key.encryption_header().to_le_bytes(),
            );
        }
        if self.snap_version.is_not_zero() {
            Builder::add_property(
                buf,
                PROP_KEY_SNAP_VERSION,
                &self.snap_version.into_inner().to_le_bytes(),
            );
        }
        let checksum = self.checksum_type.checksum(&buf[(origin_len + 4)..]);
        LittleEndian::write_u32(&mut buf[origin_len..], checksum);
    }

    pub fn is_empty(&self) -> bool {
        self.smallest.is_empty()
    }

    pub fn get_smallest(&self) -> &[u8] {
        self.smallest.as_slice()
    }

    pub fn get_biggest(&self) -> &[u8] {
        self.biggest.as_slice()
    }

    pub fn get_compression_type(&self) -> u8 {
        self.block_builder.compression_tp
    }

    pub fn get_compression_level(&self) -> i32 {
        self.block_builder.compression_lvl
    }

    pub fn get_total_blob_size(&self) -> u64 {
        self.total_blob_size
    }
}

// For test purpose.
#[cfg(test)]
impl Builder {
    fn set_old_block_uncompressed_size_limit(&mut self, limit: usize) {
        self.old_block_uncompressed_size_limit = limit;
    }
}

pub const FOOTER_SIZE: usize = mem::size_of::<Footer>();

#[repr(C)]
#[derive(Default, Clone, Copy, Debug, PartialEq, Eq)]
pub struct Footer {
    pub old_data_offset: u32,
    pub index_offset: u32,
    pub old_index_offset: u32,
    pub aux_index_offset: u32,
    pub properties_offset: u32,
    pub compression_type: u8,
    pub checksum_type: u8,
    pub table_format_version: u16,
    pub magic: u32,
}

impl Footer {
    pub fn data_len(&self) -> usize {
        self.old_data_offset as usize
    }

    pub fn old_data_len(&self) -> usize {
        (self.index_offset - self.old_data_offset) as usize
    }

    pub fn index_len(&self) -> usize {
        (self.old_index_offset - self.index_offset) as usize
    }

    pub fn old_index_len(&self) -> usize {
        (self.aux_index_offset - self.old_index_offset) as usize
    }

    pub fn aux_index_len(&self) -> usize {
        (self.properties_offset - self.aux_index_offset) as usize
    }

    pub fn properties_len(&self, table_size: usize) -> usize {
        table_size - self.properties_offset as usize - FOOTER_SIZE
    }

    /// Validates that every section offset is ordered and inside the SST.
    pub fn validate_offsets(&self, table_size: usize) -> Result<()> {
        let max_offset = table_size
            .checked_sub(FOOTER_SIZE)
            .ok_or(Error::InvalidFileSize)?;
        let offsets = [
            self.old_data_offset,
            self.index_offset,
            self.old_index_offset,
            self.aux_index_offset,
            self.properties_offset,
        ];
        if offsets.windows(2).any(|pair| pair[0] > pair[1])
            || self.properties_offset as usize > max_offset
        {
            return Err(Error::InvalidFileSize);
        }
        Ok(())
    }

    // Index is the first meta section after data sections.
    pub fn meta_offset(&self) -> u32 {
        self.index_offset
    }

    // Tiny meta is the data necessary for opening a SsTable.
    pub fn tiny_meta_offset(&self) -> u32 {
        self.properties_offset
    }

    pub fn unmarshal(&mut self, mut data: &[u8]) {
        self.old_data_offset = data.get_u32_le();
        self.index_offset = data.get_u32_le();
        self.old_index_offset = data.get_u32_le();
        self.aux_index_offset = data.get_u32_le();
        self.properties_offset = data.get_u32_le();
        self.compression_type = data.get_u8();
        self.checksum_type = data.get_u8();
        self.table_format_version = data.get_u16_le();
        self.magic = data.get_u32_le();
    }

    pub fn marshal(&self, buf: &mut Vec<u8>) {
        buf.put_u32_le(self.old_data_offset);
        buf.put_u32_le(self.index_offset);
        buf.put_u32_le(self.old_index_offset);
        buf.put_u32_le(self.aux_index_offset);
        buf.put_u32_le(self.properties_offset);
        buf.put_u8(self.compression_type);
        buf.put_u8(self.checksum_type);
        buf.put_u16_le(self.table_format_version);
        buf.put_u32_le(self.magic);
    }

    /// Note:
    ///
    /// - As L0 tables also use the same value of magic, L1+ SST tables should
    ///   be checked first when detecting SST type.
    ///
    /// - The correctness depends on the `table_format_version` field, which is
    ///   not very reliable as it's overlapped with `num_cfs` field of L0
    ///   footer.
    pub fn is_match(&self) -> bool {
        (self.magic == MAGIC_NUMBER || self.magic == MAGIC_NUMBER_SPLIT_L0)
            && self.table_format_version == TABLE_FORMAT_V1
    }
}

#[derive(Default)]
struct BlockBuffer {
    tmp_keys: EntrySlice,
    tmp_vals: EntrySlice,
    old_vers: Vec<u64>,
    entry_sizes: Vec<u32>,
    kv_size: usize,
    common_prefix_len: usize,
}

impl BlockBuffer {
    fn reset(&mut self) {
        self.tmp_keys.reset();
        self.tmp_vals.reset();
        self.old_vers.truncate(0);
        self.entry_sizes.truncate(0);
        self.kv_size = 0;
        self.common_prefix_len = 0;
    }

    fn build_entry(&self, buf: &mut Vec<u8>, i: usize, common_prefix_len: usize) {
        let key = self.tmp_keys.get_entry(i);
        let key_suffix = &key[common_prefix_len..];
        // The key suffix length is encoded as a u16. Remember to update the entry size
        // calculation (in fn add_entry()) if the key suffix length type changes.
        buf.put_u16_le(key_suffix.len() as u16);
        buf.extend_from_slice(key_suffix);
        let val_bin = self.tmp_vals.get_entry(i);
        let v = Value::decode(val_bin);
        let mut meta = v.meta;
        let old_ver = self.old_vers[i];
        if old_ver != 0 {
            meta |= BIT_HAS_OLD_VERSION;
        } else {
            // The val meta from the old table may have `metaHasOld` flag, need to unset it.
            meta &= !BIT_HAS_OLD_VERSION;
        }
        buf.push(meta);
        buf.put_u64_le(v.version);
        if old_ver != 0 {
            buf.put_u64_le(old_ver);
        }
        buf.push(v.user_meta().len() as u8);
        buf.extend_from_slice(v.user_meta());
        buf.extend_from_slice(v.get_value());
    }
}

/// Used to build data blocks from key-value entries for storage.
///
/// The Block format is as follows:
/// +---------------------+
/// | Checksum            |  (4 bytes)
/// +---------------------+
/// | Block Format ID     |  (4 bytes)
/// +---------------------+
/// | Number of Entries   |  (4 bytes)
/// +---------------------+
/// | Entry Offset 1      |  (4 bytes)
/// +---------------------+
/// | Entry Offset 2      |  (4 bytes)
/// +---------------------+
/// | ...                 |
/// +---------------------+
/// | Common Prefix Length |  (2 bytes)
/// +---------------------+
/// | Common Prefix Data   |  (variable length)
/// +---------------------+
/// | Entry 1 Key         |  (variable length)
/// | Entry 1 Value       |  (variable length)
/// +---------------------+
/// | Entry 2 Key         |  (variable length)
/// | Entry 2 Value       |  (variable length)
/// +---------------------+
/// | ...                 |
/// +---------------------+
#[derive(Default)]
pub struct BlockBuilder {
    // Final block data after encoding and compression
    buf: Vec<u8>,
    // Temporary buffer for collecting entries before building block
    block: BlockBuffer,
    // First key of each block for building index
    block_keys: EntrySlice,
    // Block addresses for tracking data locations
    block_addrs: Vec<BlockAddress>,

    // Compression type: NO_COMPRESSION(0), LZ4(1), ZSTD(2)
    compression_tp: u8,
    // Compression level for ZSTD compression
    compression_lvl: i32,
    // Temporary buffer for compression
    compression_buf: Vec<u8>,
}

impl BlockBuilder {
    pub fn get_buf(&self) -> &Vec<u8> {
        &self.buf
    }

    fn last_key_order(&self, key: &[u8]) -> Option<Ordering> {
        if self.block.tmp_keys.length() > 0 {
            let last = self.block.tmp_keys.get_last();
            return Some(last.cmp(key));
        }
        None
    }

    fn set_last_entry_old_ver_if_zero(&mut self, ver: u64) {
        let last_old_ver_idx = self.block.old_vers.len() - 1;
        if self.block.old_vers[last_old_ver_idx] == 0 {
            self.block.old_vers[last_old_ver_idx] = ver;
            let last_entry_size_idx = self.block.entry_sizes.len() - 1;
            self.block.entry_sizes[last_entry_size_idx] += VALUE_VERSION_LEN as u32;
        }
    }

    pub fn add_entry(&mut self, key: &[u8], val: Value, blob_ref: Option<BlobRef>) {
        self.block.tmp_keys.append(key);
        self.block.tmp_vals.append_value(val, blob_ref);
        self.block.old_vers.push(0);
        let encoded_size = if blob_ref.is_some() {
            val.encoded_size_with_blob_ref()
        } else {
            val.encoded_size()
        };
        let entry_size = /*key_suffix length in bytes*/ 2 + key.len() + encoded_size;
        self.block.entry_sizes.push(entry_size as u32);
        self.block.kv_size += entry_size;
        self.block.common_prefix_len = self.get_block_common_prefix_len();
    }

    fn uncompressed_block_size_with(
        num_entries: usize,
        kv_size: usize,
        common_prefix_len: usize,
    ) -> usize {
        // The uncompressed block layout (excluding checksum):
        // - format id (4)
        // - num entries (4)
        // - entry offsets (4 * n)
        // - common prefix length (2)
        // - common prefix bytes (common_prefix_len)
        // - entries data (kv_size - n * common_prefix_len)
        4 + 4
            + num_entries * 4
            + 2
            + common_prefix_len
            + (kv_size - num_entries * common_prefix_len)
    }

    fn uncompressed_block_size(&self) -> usize {
        let num_entries = self.block.tmp_keys.length();
        if num_entries == 0 {
            return 0;
        }
        Self::uncompressed_block_size_with(
            num_entries,
            self.block.kv_size,
            self.block.common_prefix_len,
        )
    }

    fn need_finish_block(&self, target_block_size: usize) -> bool {
        self.block_size() > target_block_size
    }

    fn block_size(&self) -> usize {
        self.block.kv_size - self.block.tmp_keys.length() * self.block.common_prefix_len
    }

    // Note: on error, the buffer content is unspecified/invalid.
    // Callers must reset (or drop) the builder before reusing it.
    pub fn finish_block(&mut self, sst_fid: u64, checksum_tp: ChecksumType) -> Result<()> {
        self.block_keys.append(self.block.tmp_keys.get_entry(0));
        self.block_addrs
            .push(BlockAddress::new(sst_fid, self.buf.len() as u32));
        self.buf.put_u32_le(0); // checksum place holder.
        let begin_off = self.buf.len();
        let common_prefix_len = self.get_block_common_prefix_len();
        let buf = if self.compression_tp == NO_COMPRESSION {
            &mut self.buf
        } else {
            self.compression_buf.truncate(0);
            &mut self.compression_buf
        };
        let num_entries = self.block.tmp_keys.length();
        buf.put_u32_le(BLOCK_FORMAT_V1);
        buf.put_u32_le(num_entries as u32);
        let mut offset = 0u32;
        for i in 0..num_entries {
            buf.put_u32_le(offset);
            // The entry size calculated in the first pass use full key size, we need to
            // subtract common prefix size.
            offset += self.block.entry_sizes[i] - common_prefix_len as u32;
        }
        buf.put_u16_le(common_prefix_len as u16);
        let common_prefix = &self.block.tmp_keys.get_entry(0)[..common_prefix_len];
        buf.extend_from_slice(common_prefix);
        for i in 0..num_entries {
            self.block.build_entry(buf, i, common_prefix_len);
        }
        match self.compression_tp {
            NO_COMPRESSION => (),
            LZ4_COMPRESSION => self.compress_lz4()?,
            ZSTD_COMPRESSION => self.compress_zstd()?,
            _ => panic!("unexpected compression type {}", self.compression_tp),
        }
        let checksum = checksum_tp.checksum(&self.buf[begin_off..]);
        let slice = self.buf.as_mut_slice();
        LittleEndian::write_u32(&mut slice[(begin_off - 4)..], checksum);
        self.block.reset();
        Ok(())
    }

    fn get_block_common_prefix_len(&self) -> usize {
        let first_key = self.block.tmp_keys.get_entry(0);
        let last_key = self.block.tmp_keys.get_last();
        key_diff_idx(first_key, last_key)
    }

    fn get_index_common_prefix_len(&self) -> usize {
        let first_key = self.block_keys.get_entry(0);
        let last_key = self.block_keys.get_last();
        key_diff_idx(first_key, last_key)
    }

    fn reset_all(&mut self) {
        self.block.reset();
        self.buf.truncate(0);
        self.block_keys.reset();
        self.block_addrs.truncate(0);
    }

    fn build_index(&mut self, base_off: u32, checksum_tp: ChecksumType) {
        self.buf.truncate(0);
        let num_blocks = self.block_addrs.len();
        // checksum place holder.
        self.buf.put_u32_le(0);
        self.buf.put_u32_le(INDEX_FORMAT_V1);
        self.buf.put_u32_le(num_blocks as u32);
        let mut common_prefix_len = 0;
        if num_blocks > 0 {
            common_prefix_len = self.get_index_common_prefix_len();
        }
        let mut key_offset = 0u32;
        for i in 0..num_blocks {
            self.buf.put_u32_le(key_offset);
            let block_key = self.block_keys.get_entry(i);
            key_offset += block_key.len() as u32 - common_prefix_len as u32;
        }
        for i in 0..num_blocks {
            let block_addr = &self.block_addrs[i];
            self.buf.put_u64_le(block_addr.origin_fid);
            self.buf.put_u32_le(block_addr.origin_off + base_off);
            self.buf.put_u32_le(block_addr.curr_off + base_off);
        }
        self.buf.put_u16_le(common_prefix_len as u16);
        if common_prefix_len > 0 {
            let common_prefix = &self.block_keys.get_entry(0)[..common_prefix_len];
            self.buf.extend_from_slice(common_prefix);
        }
        let block_keys_len = self.block_keys.buf.len() - num_blocks * common_prefix_len;
        self.buf.put_u32_le(block_keys_len as u32);
        for i in 0..num_blocks {
            let block_key = self.block_keys.get_entry(i);
            self.buf.extend_from_slice(&block_key[common_prefix_len..]);
        }
        let slice = self.buf.as_mut_slice();
        let checksum = checksum_tp.checksum(&slice[4..]);
        LittleEndian::write_u32(slice, checksum)
    }

    fn compress_lz4(&mut self) -> Result<()> {
        let src_len = self.compression_buf.len();
        let origin_len = self.buf.len();
        self.buf.put_u32_le(src_len as u32);
        if let Err(e) = lz4_compress_append(&self.compression_buf, &mut self.buf) {
            self.buf.truncate(origin_len);
            return Err(e);
        }
        Ok(())
    }

    fn compress_zstd(&mut self) -> Result<()> {
        let origin_len = self.buf.len();
        if let Err(e) =
            zstd_compress_append(&self.compression_buf, self.compression_lvl, &mut self.buf)
        {
            self.buf.truncate(origin_len);
            return Err(e);
        }
        Ok(())
    }
}

#[derive(Default, Clone, Copy, Debug)]
pub struct BlockAddress {
    pub origin_fid: u64,
    pub origin_off: u32,
    pub curr_off: u32,
}

impl BlockAddress {
    fn new(sst_fid: u64, offset: u32) -> Self {
        Self {
            origin_fid: sst_fid,
            origin_off: offset,
            curr_off: offset,
        }
    }

    pub(crate) fn from_slice(mut data: &[u8]) -> Self {
        let origin_fid = data.get_u64_le();
        let origin_off = data.get_u32_le();
        let curr_off = data.get_u32_le();
        Self {
            origin_fid,
            origin_off,
            curr_off,
        }
    }
}

pub struct BuildResult {
    pub id: u64,
    pub meta_offset: u32,
    pub smallest: Vec<u8>,
    pub biggest: Vec<u8>,

    #[cfg(test)]
    pub tiny_meta_offset: u32,
}

pub(crate) fn key_diff_idx(k1: &[u8], k2: &[u8]) -> usize {
    let mut i: usize = 0;
    while i < k1.len() && i < k2.len() {
        if k1[i] != k2[i] {
            break;
        }
        i += 1;
    }
    i
}

#[cfg(test)]
mod tests {
    use cloud_encryption::EncryptionState;
    use txn_types::{Lock, LockType, SharedLocks, TimeStamp};

    use super::*;
    use crate::table::{
        BIT_BLOB_REF,
        sstable::{BlockCache, BlockCacheType, SsTable},
    };

    #[test]
    fn test_entry_slice() {
        let mut es = EntrySlice::default();
        es.append("abc".as_bytes());
        let val_buf = Value::encode_buf(1, &[1], 1, "abc".as_bytes());
        let val = Value::decode(&val_buf);
        es.append_value(val, None);
        // dbg!(es.buf);
        // dbg!(es.end_offs);
    }

    #[test]
    fn test_property_key() {
        assert_eq!("max_ts".as_bytes(), b"max_ts");
    }

    // Helper function to create value for testing since Value::new_put doesn't
    // exist
    fn create_test_value(version: u64, v: &[u8]) -> Value {
        Value::new_with_version(v, version)
    }

    // Helper to create BlobRef with all required fields
    fn create_blob_ref(fid: u64, offset: u32) -> BlobRef {
        BlobRef {
            fid,
            offset,
            len: 100,          // Add required length field
            original_len: 100, // Add required original_len field
        }
    }

    #[test]
    fn test_entry_slice_append() {
        let mut slice = EntrySlice::default();
        slice.append(b"key1");
        slice.append(b"key2");

        assert_eq!(slice.length(), 2);
        assert_eq!(slice.get_entry(0), b"key1");
        assert_eq!(slice.get_entry(1), b"key2");
        assert_eq!(slice.get_last(), b"key2");
    }

    #[test]
    fn test_entry_slice_append_value() {
        let mut slice = EntrySlice::default();
        let mut value = create_test_value(123, b"test_value");

        slice.append_value(value, None);
        assert_eq!(slice.length(), 1);

        value.meta |= BIT_BLOB_REF;
        let blob_ref = create_blob_ref(100, 200);
        slice.append_value(value, Some(blob_ref));
        assert_eq!(slice.length(), 2);
    }

    #[test]
    fn test_entry_slice_get_entry() {
        let mut slice = EntrySlice::default();
        slice.append(b"key1");
        slice.append(b"key2");

        assert_eq!(slice.get_entry(0), b"key1");
        assert_eq!(slice.get_entry(1), b"key2");
    }

    #[test]
    fn test_entry_slice_size() {
        let mut slice = EntrySlice::default();
        slice.append(b"key1");
        slice.append(b"key2");

        // Size is the sum of buf size and end_offs size (4 bytes per entry)
        let expected_size = b"key1".len() + b"key2".len() + 8;
        assert_eq!(slice.size(), expected_size);
    }

    #[test]
    fn test_entry_slice_reset() {
        let mut slice = EntrySlice::default();
        slice.append(b"key1");
        slice.append(b"key2");

        slice.reset();
        assert_eq!(slice.length(), 0);
        assert!(slice.buf.is_empty());
        assert!(slice.end_offs.is_empty());
    }

    #[test]
    fn test_block_address_new() {
        let addr = BlockAddress::new(42, 1000);

        assert_eq!(addr.origin_fid, 42);
        assert_eq!(addr.origin_off, 1000);
        assert_eq!(addr.curr_off, 1000);
    }

    #[test]
    fn test_block_address_from_slice() {
        let mut data = vec![];
        data.put_u64_le(42); // origin_fid
        data.put_u32_le(1000); // origin_off
        data.put_u32_le(2000); // curr_off

        let addr = BlockAddress::from_slice(&data);

        assert_eq!(addr.origin_fid, 42);
        assert_eq!(addr.origin_off, 1000);
        assert_eq!(addr.curr_off, 2000);
    }

    // Additional tests for BlockBuffer
    #[test]
    fn test_block_buffer_default() {
        let buffer = BlockBuffer::default();

        assert_eq!(buffer.tmp_keys.length(), 0);
        assert_eq!(buffer.tmp_vals.length(), 0);
        assert!(buffer.entry_sizes.is_empty());
        assert!(buffer.old_vers.is_empty());
    }

    #[test]
    fn test_block_buffer_reset() {
        let mut buffer = BlockBuffer::default();

        // Add some test data
        buffer.tmp_keys.append(b"test_key");
        buffer.tmp_vals.append(b"test_val");
        buffer.entry_sizes.push(10);
        buffer.old_vers.push(5);

        // Reset
        buffer.reset();

        // Verify everything is cleared
        assert_eq!(buffer.tmp_keys.length(), 0);
        assert_eq!(buffer.tmp_vals.length(), 0);
        assert!(buffer.entry_sizes.is_empty());
        assert!(buffer.old_vers.is_empty());
    }

    // Helper function to create a test BlockBuilder
    fn create_test_block_builder() -> BlockBuilder {
        BlockBuilder {
            block: BlockBuffer::default(),
            block_keys: EntrySlice::default(),
            block_addrs: Vec::new(),
            buf: Vec::new(),
            compression_buf: Vec::new(),
            compression_tp: LZ4_COMPRESSION,
            compression_lvl: 3,
        }
    }

    #[test]
    fn test_block_builder_add_entry() {
        let mut builder = create_test_block_builder();
        let raw_key = format!("key{}", 0).into_bytes();
        let key = InnerKey::from_inner_buf(&raw_key);
        let value = create_test_value(100, b"test_value");

        builder.add_entry(key.deref(), value, None);

        assert_eq!(builder.block.tmp_keys.length(), 1);
        assert_eq!(builder.block.tmp_vals.length(), 1);
    }

    #[test]
    fn test_block_builder_add_entry_with_blob_ref() {
        let mut builder = create_test_block_builder();
        let raw_key = format!("key{}", 0).into_bytes();
        let key = InnerKey::from_inner_buf(&raw_key);
        let mut value = create_test_value(100, b"test_value");
        value.meta |= BIT_BLOB_REF;
        let blob_ref = create_blob_ref(42, 100);

        builder.add_entry(key.deref(), value, Some(blob_ref));

        assert_eq!(builder.block.tmp_keys.length(), 1);
        assert_eq!(builder.block.tmp_vals.length(), 1);
    }

    #[test]
    fn test_block_builder_add_entry_with_old_version() {
        let mut builder = create_test_block_builder();
        let raw_key = format!("key{}", 0).into_bytes();
        let key = InnerKey::from_inner_buf(&raw_key);
        let mut value = create_test_value(100, b"test_value");
        value.meta |= BIT_HAS_OLD_VERSION;

        builder.add_entry(key.deref(), value, None);

        assert_eq!(builder.block.tmp_keys.length(), 1);
        assert_eq!(builder.block.tmp_vals.length(), 1);
        assert_eq!(builder.block.old_vers.len(), 1);
    }

    #[test]
    fn test_block_builder_reset() {
        let mut builder = create_test_block_builder();
        let raw_key = format!("key{}", 0).into_bytes();
        let key = InnerKey::from_inner_buf(&raw_key);
        let value = create_test_value(100, b"test_value");

        builder.add_entry(key.deref(), value, None);

        builder.reset_all();

        assert_eq!(builder.block.tmp_keys.length(), 0);
        assert_eq!(builder.block.tmp_vals.length(), 0);
        assert!(builder.block.entry_sizes.is_empty());
        assert!(builder.block.old_vers.is_empty());
        assert!(builder.buf.is_empty());
        assert_eq!(builder.block_keys.length(), 0);
        assert!(builder.block_addrs.is_empty());
    }

    #[test]
    fn test_block_builder_size() {
        let mut builder = create_test_block_builder();

        // Add some entries
        for i in 0..5 {
            let raw_key = format!("key{}", i).into_bytes();
            let key = InnerKey::from_inner_buf(&raw_key);
            let value = create_test_value(100, format!("value{}", i).as_bytes());
            builder.add_entry(key.deref(), value, None);
        }

        let size = builder.block_size();
        assert!(size > 0);
    }

    #[test]
    fn test_block_builder_finish_block() {
        let mut builder = create_test_block_builder();

        // Add some entries
        for i in 0..5 {
            let raw_key = format!("key{}", i).into_bytes();
            let key = InnerKey::from_inner_buf(&raw_key);
            let value = create_test_value(100, format!("value{}", i).as_bytes());
            builder.add_entry(key.deref(), value, None);
        }

        let sst_fid = 123;
        let checksum_type = ChecksumType::Crc32c;

        builder.finish_block(sst_fid, checksum_type).unwrap();

        // Check if the block has been reset
        assert_eq!(builder.block.tmp_keys.length(), 0);
        assert_eq!(builder.block.tmp_vals.length(), 0);
        assert!(builder.block.entry_sizes.is_empty());

        // Check if block keys and addresses have been updated
        assert_eq!(builder.block_keys.length(), 1);
        assert_eq!(builder.block_addrs.len(), 1);
        assert_eq!(builder.block_addrs[0].origin_fid, sst_fid);
    }

    #[test]
    fn test_block_builder_build_index() {
        let mut builder = create_test_block_builder();

        // Add some blocks
        for i in 0..3 {
            // Add some entries to each block
            for j in 0..3 {
                let raw_key = format!("block{}_key{}", i, j).into_bytes();
                let key = InnerKey::from_inner_buf(&raw_key);
                let value = create_test_value(100, format!("value{}_{}", i, j).as_bytes());
                builder.add_entry(key.deref(), value, None);
            }

            // Finish the block
            builder.finish_block(123, ChecksumType::Crc32c).unwrap();
        }

        // Clear the buffer before building the index
        builder.buf.clear();

        // Build the index
        let base_off = 1000;
        builder.build_index(base_off, ChecksumType::Crc32c);

        // Check if the buffer contains the index
        assert!(!builder.buf.is_empty());
    }

    #[test]
    fn test_block_builder_get_block_common_prefix_len() {
        let mut builder = create_test_block_builder();

        // Add entries with common prefix
        builder.block.tmp_keys.append(b"prefix_key1");
        builder.block.tmp_keys.append(b"prefix_key2");

        let common_prefix_len = builder.get_block_common_prefix_len();
        assert_eq!(common_prefix_len, 10); // "prefix_key" is 10 chars
    }

    #[test]
    fn test_block_builder_get_index_common_prefix_len() {
        let mut builder = create_test_block_builder();

        // Add block keys with common prefix
        builder.block_keys.append(b"prefix_block1");
        builder.block_keys.append(b"prefix_block2");

        let common_prefix_len = builder.get_index_common_prefix_len();
        assert_eq!(common_prefix_len, 12); // "prefix_block" is 12 chars
    }

    #[test]
    fn test_block_builder_compress_lz4() {
        let mut builder = create_test_block_builder();

        // Fill compression buffer with some data
        builder.compression_buf.extend_from_slice(&[0; 1000]);

        let original_compression_buf_len = builder.compression_buf.len();
        let original_buf_len = builder.buf.len();

        builder.compress_lz4().unwrap();

        // Buffer should now contain the compressed data
        assert!(builder.buf.len() > original_buf_len);
        assert!(builder.buf.len() < original_compression_buf_len);
    }

    #[test]
    fn test_block_builder_compress_zstd() {
        let mut builder = create_test_block_builder();
        builder.compression_tp = ZSTD_COMPRESSION;

        // Fill compression buffer with some data
        builder.compression_buf.extend_from_slice(&[0; 1000]);

        let original_buf_len = builder.buf.len();
        let original_compression_buf_len = builder.compression_buf.len();

        builder.compress_zstd().unwrap();

        // Buffer should now contain the compressed data
        assert!(builder.buf.len() > original_buf_len);
        assert!(builder.buf.len() < original_compression_buf_len);
    }

    #[test]
    fn test_block_builder_set_last_entry_old_ver_if_zero() {
        let mut builder = create_test_block_builder();

        // Add an entry to the builder
        let raw_key = b"test_key".to_vec();
        let key = InnerKey::from_inner_buf(&raw_key);
        let value = create_test_value(100, b"test_value");
        builder.add_entry(key.deref(), value, None);

        // Initially, the old version should be zero
        assert_eq!(builder.block.old_vers.len(), 1);
        assert_eq!(builder.block.old_vers[0], 0);

        // Call the method with a version
        builder.set_last_entry_old_ver_if_zero(42);

        // Verify that the old version has been set
        assert_eq!(builder.block.old_vers[0], 42);

        // Call the method again with a different version
        builder.set_last_entry_old_ver_if_zero(100);

        // Verify that the old version has not changed
        assert_eq!(builder.block.old_vers[0], 42); // Should still be 42
    }

    #[test]
    fn test_block_builder_set_last_entry_old_ver_if_non_zero() {
        let mut builder = create_test_block_builder();

        // Add an entry to the builder
        let raw_key = b"test_key".to_vec();
        let key = InnerKey::from_inner_buf(&raw_key);
        let value = create_test_value(100, b"test_value");
        builder.add_entry(key.deref(), value, None);

        // Set the old version manually to a non-zero value
        builder.block.old_vers[0] = 10;

        // Call the method with a new version
        builder.set_last_entry_old_ver_if_zero(42);

        // Verify that the old version has not changed
        assert_eq!(builder.block.old_vers[0], 10); // Should still be 10
    }

    #[test]
    fn test_table_builder_options_default() {
        let opts = TableBuilderOptions::default();

        assert_eq!(opts.block_size, 64 * 1024);
        assert_eq!(opts.max_table_size, 16 * 1024 * 1024);
        assert_eq!(opts.compression_tps[0], LZ4_COMPRESSION);
        assert_eq!(opts.compression_tps[1], ZSTD_COMPRESSION);
        assert_eq!(opts.compression_tps[2], ZSTD_COMPRESSION);
        assert_eq!(opts.compression_lvl, 3);
        assert_eq!(opts.flush_split_l0, false);
    }

    // Helper function to create a test Builder instance
    fn create_test_builder() -> Builder {
        Builder::new(
            123,                  // sst_fid
            4096,                 // block_size
            LZ4_COMPRESSION,      // compression_tp
            3,                    // compression_lvl
            ChecksumType::Crc32c, // checksum_type
            None,                 // encryption_key
        )
    }

    // Helper function to create a test encryption key
    fn create_test_encryption_key() -> Option<EncryptionKey> {
        // Since the actual structure is different than what we tried to test,
        // let's just return None for now
        None
    }

    #[test]
    fn test_builder_new() {
        let builder = create_test_builder();

        assert_eq!(builder.sst_fid, 123);
        assert_eq!(builder.block_size, 4096);
        assert_eq!(builder.checksum_type, ChecksumType::Crc32c);
        assert!(builder.smallest.is_empty());
        assert!(builder.biggest.is_empty());
        assert_eq!(builder.max_ts, 0);
        assert_eq!(builder.old_entries, 0);
        assert_eq!(builder.tombs, 0);
        assert_eq!(builder.kv_size, 0);
    }

    #[test]
    fn test_builder_reset() {
        let mut builder = create_test_builder();

        // Add some data
        builder.smallest = b"small".to_vec();
        builder.biggest = b"big".to_vec();
        builder.max_ts = 100;
        builder.old_entries = 5;
        builder.tombs = 3;
        builder.kv_size = 1000;

        // Reset the builder
        builder.reset(456);

        assert_eq!(builder.sst_fid, 456);
        assert!(builder.smallest.is_empty());
        assert!(builder.biggest.is_empty());
        assert_eq!(builder.max_ts, 0);
        assert_eq!(builder.old_entries, 0);
        assert_eq!(builder.tombs, 0);
        assert_eq!(builder.kv_size, 0);
    }

    #[test]
    fn test_builder_tombs_count_main_region_only() {
        // A tombstone in the latest (main-region) version of a key is counted,
        // but a tombstone that is only an old version of a key is not. This
        // keeps `tombs` consistent with `entries` (which only covers the main
        // region). See issue #5409.

        // Case 1: latest version is a tombstone -> counted.
        let mut builder = create_test_builder();
        let raw_key = b"key1".to_vec();
        let key = InnerKey::from_inner_buf(&raw_key);
        // Higher version is added first (descending order within a key).
        let tomb = Value::new_tombstone(100);
        builder.add(key, &tomb, None).unwrap();
        assert_eq!(builder.tombs, 1);
        assert_eq!(builder.old_entries, 0);

        // Case 2: latest version has a value, old version is a tombstone ->
        // only the old version is a tombstone, so it must NOT be counted.
        let mut builder = create_test_builder();
        let raw_key = b"key1".to_vec();
        let key = InnerKey::from_inner_buf(&raw_key);
        let latest = create_test_value(100, b"value");
        builder.add(key, &latest, None).unwrap();
        let old_tomb = Value::new_tombstone(50);
        builder.add(key, &old_tomb, None).unwrap();
        assert_eq!(builder.old_entries, 1);
        assert_eq!(builder.tombs, 0);

        let mut buf = Vec::with_capacity(builder.estimated_size());
        builder.finish(0, &mut buf).unwrap();
        let table = SsTable::new(
            std::sync::Arc::new(crate::table::file::InMemFile::new(123, buf.into())),
            BlockCache::new(BlockCacheType::None, 0, 4096),
            EncryptionState::Disabled,
        )
        .unwrap();
        assert_eq!(table.entries, 1);
        assert_eq!(table.old_entries, 1);
        assert_eq!(table.tombs, 0);
    }

    #[test]
    fn test_builder_is_empty() {
        let builder = Builder::new(
            123,                  // sst_fid
            4096,                 // block_size
            LZ4_COMPRESSION,      // compression_tp
            3,                    // compression_lvl
            ChecksumType::Crc32c, // checksum_type
            None,                 // encryption_key
        );

        assert!(builder.is_empty()); // Should be true for a new builder

        let raw_key = b"test_key".to_vec();
        let key = InnerKey::from_inner_buf(&raw_key);
        let value = create_test_value(100, b"test_value");

        let mut builder_with_data = builder;
        builder_with_data.add(key, &value, None).unwrap();

        assert!(!builder_with_data.is_empty()); // Should be false after adding an entry
    }

    #[test]
    fn test_builder_smallest_and_biggest() {
        let mut builder = Builder::new(
            123,                  // sst_fid
            4096,                 // block_size
            LZ4_COMPRESSION,      // compression_tp
            3,                    // compression_lvl
            ChecksumType::Crc32c, // checksum_type
            None,                 // encryption_key
        );

        // Insert multiple keys
        builder
            .add(
                InnerKey::from_inner_buf(b"key1"),
                &create_test_value(100, b"value1"),
                None,
            )
            .unwrap();
        builder
            .add(
                InnerKey::from_inner_buf(b"key2"),
                &create_test_value(100, b"value2"),
                None,
            )
            .unwrap();
        builder
            .add(
                InnerKey::from_inner_buf(b"key3"),
                &create_test_value(100, b"value3"),
                None,
            )
            .unwrap();
        let mut buf = Vec::new();
        let result = builder.finish(0, &mut buf).unwrap();

        assert_eq!(result.smallest, b"key1");
        assert_eq!(result.biggest, b"key3");
    }

    #[test]
    fn test_lock_cf_max_lock_ts_ignores_invalid_lock_bytes() {
        let mut builder = create_test_builder();
        builder.set_is_lock_cf();
        let value = Value::new_with_meta_version(0, 1, 0, b"not-a-lock");
        builder
            .add(InnerKey::from_inner_buf(b"key"), &value, None)
            .unwrap();

        assert_eq!(builder.max_lock_ts, 0);
    }

    #[test]
    fn test_lock_cf_max_lock_ts_counts_shared_lock_holders() {
        let mut shared_locks = SharedLocks::new();
        for ts in [20_u64, 40, 30] {
            shared_locks
                .insert_lock(Lock::new(
                    LockType::Lock,
                    format!("pk-{ts}").into_bytes(),
                    ts.into(),
                    100,
                    None,
                    TimeStamp::zero(),
                    0,
                    TimeStamp::zero(),
                ))
                .unwrap();
        }

        let mut builder = create_test_builder();
        builder.set_is_lock_cf();
        let shared_lock_bytes = shared_locks.to_bytes();
        let value = Value::new_with_meta_version(0, 1, 0, &shared_lock_bytes);
        builder
            .add(InnerKey::from_inner_buf(b"key"), &value, None)
            .unwrap();

        assert_eq!(builder.max_lock_ts, 40);
    }

    #[test]
    fn test_builder_estimated_size() {
        let mut builder = create_test_builder();

        // Add some data to increase the estimated size
        for i in 0..10 {
            let raw_key = format!("key{}", i).into_bytes();
            let key = InnerKey::from_inner_buf(&raw_key);
            let value = create_test_value(100, format!("value{}", i).as_bytes());
            builder.add(key, &value, None).unwrap();
        }

        let size = builder.estimated_size();
        assert!(size > 0);
    }

    #[test]
    fn test_builder_finish() {
        let mut builder = create_test_builder();

        // Add some data
        for i in 0..10 {
            let raw_key = format!("key{}", i).into_bytes();
            let key = InnerKey::from_inner_buf(&raw_key);
            let value = create_test_value(100, format!("value{}", i).as_bytes());
            builder.add(key, &value, None).unwrap();
        }

        let mut buf = Vec::new();
        let result = builder.finish(0, &mut buf).unwrap();

        assert_eq!(result.id, 123);
        assert!(!buf.is_empty());
    }

    #[test]
    fn test_old_block_uncompressed_size_limit_exceeded() {
        let key = InnerKey::from_inner_buf(b"key");
        let v3 = create_test_value(3, b"value");
        let v2 = create_test_value(2, b"value");
        let v1 = create_test_value(1, b"value");

        // Probe the exact uncompressed old-block sizes for 1 and 2 old entries.
        let mut probe = Builder::new(
            123,                  // sst_fid
            4096,                 // block_size
            LZ4_COMPRESSION,      // compression_tp
            3,                    // compression_lvl
            ChecksumType::Crc32c, // checksum_type
            None,                 // encryption_key
        );

        probe.set_old_block_uncompressed_size_limit(usize::MAX);
        probe.add(key, &v3, None).unwrap();
        probe.add(key, &v2, None).unwrap();
        let size_for_one_old = probe.old_builder.uncompressed_block_size();
        probe.add(key, &v1, None).unwrap();
        let size_for_two_old = probe.old_builder.uncompressed_block_size();
        assert!(size_for_two_old > size_for_one_old);

        let mut builder = Builder::new(
            123,                  // sst_fid
            4096,                 // block_size
            LZ4_COMPRESSION,      // compression_tp
            3,                    // compression_lvl
            ChecksumType::Crc32c, // checksum_type
            None,                 // encryption_key
        );

        // Latest version goes to the data block builder.
        builder.add(key, &v3, None).unwrap();

        // Old versions are appended into the old-version block builder.
        builder.set_old_block_uncompressed_size_limit(size_for_one_old);
        builder.add(key, &v2, None).unwrap();

        builder.set_old_block_uncompressed_size_limit(size_for_two_old - 1);
        let err = builder.add(key, &v1, None).unwrap_err();
        match err {
            Error::OldVersionBlockTooLarge { size, limit } => {
                assert_eq!(size, size_for_two_old);
                assert_eq!(limit, size_for_two_old - 1);
            }
            other => panic!("unexpected error: {:?}", other),
        }
    }

    #[test]
    fn test_builder_add_with_keyspace_id() {
        let mut builder = Builder::new(
            123,                  // sst_fid
            4096,                 // block_size
            LZ4_COMPRESSION,      // compression_tp
            3,                    // compression_lvl
            ChecksumType::Crc32c, // checksum_type
            None,                 // encryption_key
        );

        let raw_key = b"test_key".to_vec();
        let key = InnerKey::from_inner_buf(&raw_key);
        let value = create_test_value(100, b"test_value");

        builder.add(key, &value, None).unwrap();

        // Should have successfully added entry with keyspace ID prepended
        assert!(!builder.block_builder.block.tmp_keys.buf.is_empty());
    }

    #[test]
    fn test_builder_add_out_of_order() {
        let mut builder = create_test_builder();

        let raw_key = b"test_key_1".to_vec();
        let key = InnerKey::from_inner_buf(&raw_key);
        let value = create_test_value(100, b"test_value");

        builder.add(key, &value, None).unwrap();

        let out_of_order_key = b"test_key_0".to_vec();
        let out_of_order_key = InnerKey::from_inner_buf(&out_of_order_key);
        assert!(builder.add(out_of_order_key, &value, None).is_err());
    }

    #[test]
    fn test_builder_add_with_blob_ref() {
        let mut builder = create_test_builder();

        let raw_key = b"test_key".to_vec();
        let key = InnerKey::from_inner_buf(&raw_key);
        let mut value = create_test_value(100, b"test_value");
        value.meta |= BIT_BLOB_REF;

        let blob_ref = create_blob_ref(42, 100);

        builder.add(key, &value, Some(blob_ref)).unwrap();

        // Should have successfully added entry with blob reference
        assert!(!builder.block_builder.block.tmp_keys.buf.is_empty());
    }

    // Test for Builder's add method with encrypted data
    #[test]
    fn test_builder_with_encryption() {
        // This test depends on whether we can create a mock encryption key
        // For now, just verify the code path works without actual encryption
        let encryption_key = create_test_encryption_key();

        let mut builder = Builder::new(
            123,                  // sst_fid
            4096,                 // block_size
            LZ4_COMPRESSION,      // compression_tp
            3,                    // compression_lvl
            ChecksumType::Crc32c, // checksum_type
            encryption_key,       // encryption_key (None for now)
        );

        let raw_key = b"test_key".to_vec();
        let key = InnerKey::from_inner_buf(&raw_key);
        let value = create_test_value(100, b"test_value");

        builder.add(key, &value, None).unwrap();

        // Should have successfully added entry
        assert!(!builder.block_builder.block.tmp_keys.buf.is_empty());
    }

    #[test]
    fn test_footer_data_len() {
        let mut footer = Footer::default();
        footer.old_data_offset = 1000;

        assert_eq!(footer.data_len(), 1000);
    }

    #[test]
    fn test_footer_old_data_len() {
        let mut footer = Footer::default();
        footer.old_data_offset = 1000;
        footer.index_offset = 1500;

        assert_eq!(footer.old_data_len(), 500); // 1500 - 1000 = 500
    }

    #[test]
    fn test_footer_index_len() {
        let mut footer = Footer::default();
        footer.index_offset = 1500;
        footer.old_index_offset = 2000;

        assert_eq!(footer.index_len(), 500); // 2000 - 1500 = 500
    }

    #[test]
    fn test_footer_old_index_len() {
        let mut footer = Footer::default();
        footer.old_index_offset = 2000;
        footer.aux_index_offset = 2500;

        assert_eq!(footer.old_index_len(), 500); // 2500 - 2000 = 500
    }

    #[test]
    fn test_footer_aux_index_len() {
        let mut footer = Footer::default();
        footer.aux_index_offset = 2500;
        footer.properties_offset = 3000;

        assert_eq!(footer.aux_index_len(), 500); // 3000 - 2500 = 500
    }

    #[test]
    fn test_footer_properties_len() {
        let mut footer = Footer::default();
        footer.properties_offset = 3000;

        // Assuming the total size of the table is 4000 for this test
        let table_size = 4000;

        // Correctly calculate the properties length considering FOOTER_SIZE
        assert_eq!(
            footer.properties_len(table_size),
            table_size - footer.properties_offset as usize - FOOTER_SIZE
        );
    }

    #[test]
    fn test_footer_marshal() {
        let mut footer = Footer::default();
        footer.old_data_offset = 100;
        footer.index_offset = 200;
        footer.old_index_offset = 300;
        footer.aux_index_offset = 400;
        footer.properties_offset = 500;
        footer.compression_type = LZ4_COMPRESSION;
        footer.checksum_type = ChecksumType::Crc32c.value();
        footer.table_format_version = TABLE_FORMAT_V1;
        footer.magic = MAGIC_NUMBER;

        let mut buf = Vec::new();
        footer.marshal(&mut buf);

        // Verify buffer size is correct
        assert_eq!(buf.len(), FOOTER_SIZE);

        // Verify magic number is at the end (last 4 bytes)
        let magic_bytes = &buf[buf.len() - 4..];
        let magic = LittleEndian::read_u32(magic_bytes);
        assert_eq!(magic, MAGIC_NUMBER);
    }

    #[test]
    fn test_footer_unmarshal() {
        let mut original_footer = Footer::default();
        original_footer.old_data_offset = 100;
        original_footer.index_offset = 200;
        original_footer.old_index_offset = 300;
        original_footer.aux_index_offset = 400;
        original_footer.properties_offset = 500;
        original_footer.compression_type = LZ4_COMPRESSION;
        original_footer.checksum_type = ChecksumType::Crc32c.value();
        original_footer.table_format_version = TABLE_FORMAT_V1;
        original_footer.magic = MAGIC_NUMBER;

        let mut buf = Vec::new();
        original_footer.marshal(&mut buf);

        // Now unmarshal back into a new footer
        let mut new_footer = Footer::default();
        new_footer.unmarshal(&buf);

        // Verify all fields match
        assert_eq!(new_footer.old_data_offset, original_footer.old_data_offset);
        assert_eq!(new_footer.index_offset, original_footer.index_offset);
        assert_eq!(
            new_footer.old_index_offset,
            original_footer.old_index_offset
        );
        assert_eq!(
            new_footer.aux_index_offset,
            original_footer.aux_index_offset
        );
        assert_eq!(
            new_footer.properties_offset,
            original_footer.properties_offset
        );
        assert_eq!(
            new_footer.compression_type,
            original_footer.compression_type
        );
        assert_eq!(new_footer.checksum_type, original_footer.checksum_type);
        assert_eq!(
            new_footer.table_format_version,
            original_footer.table_format_version
        );
        assert_eq!(new_footer.magic, original_footer.magic);
    }

    #[test]
    fn test_footer_is_match_valid() {
        let mut footer = Footer::default();
        footer.magic = MAGIC_NUMBER; // Valid magic number
        footer.table_format_version = TABLE_FORMAT_V1; // Valid format version

        assert!(footer.is_match());
    }

    #[test]
    fn test_footer_is_match_invalid_magic() {
        let mut footer = Footer::default();
        footer.magic = 123456; // Invalid magic number
        footer.table_format_version = TABLE_FORMAT_V1; // Valid format version

        assert!(!footer.is_match());
    }

    #[test]
    fn test_footer_is_match_invalid_format_version() {
        let mut footer = Footer::default();
        footer.magic = MAGIC_NUMBER; // Valid magic number
        footer.table_format_version = 999; // Invalid format version

        assert!(!footer.is_match());
    }

    #[test]
    fn test_footer_is_match_invalid_both() {
        let mut footer = Footer::default();
        footer.magic = 123456; // Invalid magic number
        footer.table_format_version = 999; // Invalid format version

        assert!(!footer.is_match());
    }

    #[test]
    fn test_key_diff_idx() {
        assert_eq!(key_diff_idx(b"abcdef", b"abcxyz"), 3);
        assert_eq!(key_diff_idx(b"abc", b"abc"), 3);
        assert_eq!(key_diff_idx(b"abc", b"abcdef"), 3);
        assert_eq!(key_diff_idx(b"abcdef", b"abc"), 3);
        assert_eq!(key_diff_idx(b"", b"abc"), 0);
        assert_eq!(key_diff_idx(b"abc", b""), 0);
        assert_eq!(key_diff_idx(b"", b""), 0);
        assert_eq!(key_diff_idx(b"xyz", b"abc"), 0);
    }

    #[test]
    fn test_key_diff_idx_with_edge_cases() {
        // Test with non-ASCII characters
        assert_eq!(key_diff_idx("你好".as_bytes(), "你们好".as_bytes()), 3);

        // Test with very long strings
        let long_str1 = "a".repeat(1000);
        let long_str2 = format!("{}b", "a".repeat(999));
        assert_eq!(
            key_diff_idx(long_str1.as_bytes(), long_str2.as_bytes()),
            999
        );

        // Test with binary data
        assert_eq!(key_diff_idx(&[0, 1, 2, 3], &[0, 1, 2, 4]), 3);
    }
}
