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

use std::hash::Hasher;

use twox_hash::{XxHash32, XxHash64};

use crate::{Error, Result};

const CRC_MASK_DELTA: u32 = 0xa282_ead8;
const LAST_BYTE_PRIME: u32 = 0x6b90_83d9;

/// Checksum algorithms stored in a RocksDB block-based table footer.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[repr(u8)]
pub enum ChecksumType {
    None = 0,
    Crc32c = 1,
    XxHash = 2,
    XxHash64 = 3,
    Xxh3 = 4,
}

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

    fn try_from(value: u8) -> Result<Self> {
        match value {
            0 => Ok(Self::None),
            1 => Ok(Self::Crc32c),
            2 => Ok(Self::XxHash),
            3 => Ok(Self::XxHash64),
            4 => Ok(Self::Xxh3),
            other => Err(Error::Unsupported(format!("checksum type {other}"))),
        }
    }
}

pub(crate) fn checksum(ty: ChecksumType, data: &[u8], last_byte: u8) -> u32 {
    match ty {
        ChecksumType::None => 0,
        ChecksumType::Crc32c => {
            mask_crc32c(crc32c::crc32c_append(crc32c::crc32c(data), &[last_byte]))
        }
        ChecksumType::XxHash => {
            let mut hasher = XxHash32::with_seed(0);
            hasher.write(data);
            hasher.write_u8(last_byte);
            hasher.finish() as u32
        }
        ChecksumType::XxHash64 => {
            let mut hasher = XxHash64::with_seed(0);
            hasher.write(data);
            hasher.write_u8(last_byte);
            hasher.finish() as u32
        }
        ChecksumType::Xxh3 => {
            let value = xxhash_rust::xxh3::xxh3_64(data) as u32;
            value ^ u32::from(last_byte).wrapping_mul(LAST_BYTE_PRIME)
        }
    }
}

fn mask_crc32c(crc: u32) -> u32 {
    crc.rotate_right(15).wrapping_add(CRC_MASK_DELTA)
}

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

    // Values copied from RocksDB 8.10's BuiltinChecksumTest.ChecksumSchemas.
    #[test]
    fn matches_rocksdb_checksum_schema() {
        let data = b"This is a short block!";
        assert_eq!(checksum(ChecksumType::None, data, 0), 0x0000_0000);
        assert_eq!(checksum(ChecksumType::Crc32c, data, 0), 0x5503_3f58);
        assert_eq!(checksum(ChecksumType::XxHash, data, 0), 0xb05f_2e4a);
        assert_eq!(checksum(ChecksumType::XxHash64, data, 0), 0xef55_46b7);
        assert_eq!(checksum(ChecksumType::Xxh3, data, 0), 0xe6b5_7fb3);
    }
}