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

use std::{ops::Deref, sync::Arc};

use byteorder::{ByteOrder, LittleEndian};
use bytes::{Buf, BufMut, Bytes, BytesMut};
use cloud_encryption::{EncryptionKey, EncryptionState};
use kvenginepb::L0Create;

use super::*;
use crate::{
    NUM_CFS, WRITE_CF, max_ts_by_cf,
    table::{
        BoundedDataSet, ChecksumType, DataBound, Error, InnerKey, NO_COMPRESSION, SnapVersion,
        Value, blobtable::BlobRef, file::File, table::Result,
    },
    util::new_l0_create_pb,
};

const L0_FOOTER_SIZE: usize = std::mem::size_of::<L0Footer>();
const L0_CF_OFFSETS_SIZE: usize = std::mem::size_of::<u32>() * NUM_CFS;

#[derive(Default, Clone)]
pub struct L0Footer {
    snap_version: SnapVersion,
    num_cfs: u32,
    magic: u32,
}

impl L0Footer {
    pub fn unmarshal(&mut self, bin: &[u8]) {
        self.snap_version = LittleEndian::read_u64(bin).into();
        self.num_cfs = LittleEndian::read_u32(&bin[8..]);
        self.magic = LittleEndian::read_u32(&bin[12..]);
    }

    pub fn is_match(&self) -> bool {
        self.magic == MAGIC_NUMBER || self.magic == MAGIC_NUMBER_SPLIT_L0
    }
}

// The L0Table format is as follows:
// +-------------------------+
// |         L0 File         |
// |                         |
// |  +-------------------+  |
// |  |     SSTable 1     |  | <--- Data from Builder 1 (CF 1)
// |  | +---------------+ |  |
// |  | |    Block 1    | |  |
// |  | +---------------+ |  |
// |  | |    Block 2    | |  |
// |  | +---------------+ |  |
// |  | |     ...       | |  |
// |  | +---------------+ |  |
// |  | |    Index      | |  |
// |  | +---------------+ |  |
// |  | |    Footer     | |  |
// |  | +---------------+ |  |
// |  +-------------------+  |
// |  |     SSTable 2     |  | <--- Data from Builder 2 (CF 2)
// |  |                   |  |
// |  |                   |  |
// |  +-------------------+  |
// |  |     SSTable 3     |  | <--- Data from Builder 3 (CF 3)
// |  |                   |  |
// |  |                   |  |
// |  +-------------------+  |
// |                         |
// |          ...            |
// |                         |
// |  +-------------------+  |
// |  |      Footer       |  | <--- L0 Table Footer
// |  | +-------------+ | |  |
// |  | |    Offsets    | |  | <--- Offsets of SSTables
// |  | +-------------+ | |  |
// |  | |    Version    | |  | <--- Version number
// |  | +-------------+ | |  |
// |  | |  Num of CFS   | |  | <--- Number of column families
// |  | +-------------+ | |  |
// |  | | MAGIC_NUMBER  | |  | <--- Magic number for validation
// |  | +-------------+ | |  |
// |  +-----------------+ |  |
// +-------------------------+
#[derive(Clone)]
pub struct L0Table {
    core: Arc<L0TableCore>,
}

impl Deref for L0Table {
    type Target = L0TableCore;

    fn deref(&self) -> &Self::Target {
        &self.core
    }
}

impl L0Table {
    /// Would return `None` only when `write_cf_only` is `true`.
    pub fn new(
        file: Arc<dyn File>,
        cache: BlockCache,
        write_cf_only: bool,
        encryption_state: EncryptionState,
    ) -> Result<Option<Self>> {
        let core = L0TableCore::new(file, cache, write_cf_only, encryption_state)?;
        Ok(core.map(|core| Self {
            core: Arc::new(core),
        }))
    }

    pub const fn footer_size() -> usize {
        L0_FOOTER_SIZE
    }
}

pub struct L0TableCore {
    footer: L0Footer,
    file: Arc<dyn File>,
    cfs: [Option<sstable::SsTable>; NUM_CFS],
    max_ts: u64,
    entries: u64,
    kv_size: u64,
    smallest: Bytes,
    biggest: Bytes,
    total_blob_size: u64,
}

impl L0TableCore {
    pub fn new(
        file: Arc<dyn File>,
        cache: BlockCache,
        write_cf_only: bool,
        encryption_state: EncryptionState,
    ) -> Result<Option<Self>> {
        let footer_off = file
            .size()
            .checked_sub(L0_FOOTER_SIZE as u64)
            .ok_or(Error::InvalidFileSize)?;
        let mut footer = L0Footer::default();
        let footer_buf = file.read_footer(L0_FOOTER_SIZE)?;
        footer.unmarshal(footer_buf.chunk());
        if footer.magic == MAGIC_NUMBER_SPLIT_L0 {
            return Self::new_write_cf_l0(file, cache, encryption_state);
        } else if footer.magic != MAGIC_NUMBER {
            return Err(Error::InvalidMagicNumber);
        }
        if footer.num_cfs != NUM_CFS as u32 {
            return Err(Error::InvalidFileSize);
        }
        let cf_offs_off = footer_off
            .checked_sub(L0_CF_OFFSETS_SIZE as u64)
            .ok_or(Error::InvalidFileSize)?;
        let cf_offs_buf = file.read(cf_offs_off, L0_CF_OFFSETS_SIZE)?;
        let mut cf_offs = [0u64; NUM_CFS];
        for i in 0..NUM_CFS {
            let offset = LittleEndian::read_u32(&cf_offs_buf[i * 4..]) as u64;
            if (i == 0 && offset != 0) || offset > cf_offs_off || (i > 0 && offset < cf_offs[i - 1])
            {
                return Err(Error::InvalidFileSize);
            }
            cf_offs[i] = offset;
        }
        let mut cfs: [Option<SsTable>; NUM_CFS] = [None, None, None];
        let mut entries = 0;
        let mut kv_size = 0;
        for i in 0..NUM_CFS {
            let start_off = cf_offs[i];
            let mut end_off = cf_offs_off;
            if i + 1 < NUM_CFS {
                end_off = cf_offs[i + 1];
            }
            if start_off == end_off || write_cf_only && i != WRITE_CF {
                continue;
            }
            let tbl = sstable::SsTable::new_l0_cf(
                file.clone(),
                start_off,
                end_off,
                cache.clone(),
                encryption_state.clone(),
            )?;
            entries += tbl.entries as u64;
            if i == WRITE_CF {
                kv_size += tbl.kv_size;
            }

            cfs[i] = Some(tbl)
        }

        if cfs.iter().all(|t| t.is_none()) {
            // All CFs are empty only when `write_cf_only` is `true` and `WRITE_CF` is
            // empty.
            return if write_cf_only {
                Ok(None)
            } else {
                Err(Error::InvalidFileSize)
            };
        }

        let (smallest, biggest, max_ts) = Self::compute_smallest_biggest(&cfs);
        let total_blob_size = Self::compute_total_blob_size(&cfs);
        Ok(Some(Self {
            footer,
            file,
            cfs,
            max_ts,
            entries,
            kv_size,
            smallest,
            biggest,
            total_blob_size,
        }))
    }

    fn new_write_cf_l0(
        file: Arc<dyn File>,
        cache: BlockCache,
        encryption_state: EncryptionState,
    ) -> Result<Option<Self>> {
        let tbl = SsTable::new(file.clone(), cache, encryption_state)?;
        let entries = tbl.entries as u64;
        let kv_size = tbl.kv_size;
        let mut footer = L0Footer::default();
        footer.num_cfs = NUM_CFS as u32;
        footer.magic = MAGIC_NUMBER_SPLIT_L0;
        footer.snap_version = tbl.snap_version;
        let smallest = tbl.clone_smallest();
        let biggest = tbl.clone_biggest();
        let max_ts = tbl.max_ts;
        let total_blob_size = tbl.total_blob_size();
        Ok(Some(Self {
            footer,
            file,
            cfs: [Some(tbl), None, None],
            max_ts,
            entries,
            kv_size,
            smallest,
            biggest,
            total_blob_size,
        }))
    }

    // Return: smallest, biggest, max_ts
    fn compute_smallest_biggest(cfs: &[Option<SsTable>; NUM_CFS]) -> (Bytes, Bytes, u64) {
        let mut smallest_buf = BytesMut::new();
        let mut biggest_buf = BytesMut::new();
        let mut max_ts = 0;
        for i in 0..NUM_CFS {
            if let Some(cf_tbl) = &cfs[i] {
                let smallest = cf_tbl.smallest();
                if !smallest.is_empty()
                    && (smallest_buf.is_empty() || smallest_buf.chunk() > smallest.deref())
                {
                    smallest_buf.truncate(0);
                    smallest_buf.extend_from_slice(smallest.deref());
                }
                let biggest = cf_tbl.biggest();
                if biggest.deref() > biggest_buf.chunk() {
                    biggest_buf.truncate(0);
                    biggest_buf.extend_from_slice(biggest.deref());
                }
                max_ts = max_ts_by_cf(max_ts, i, cf_tbl.max_ts);
            }
        }
        assert!(!smallest_buf.is_empty());
        assert!(!biggest_buf.is_empty());
        (smallest_buf.freeze(), biggest_buf.freeze(), max_ts)
    }

    fn compute_total_blob_size(cfs: &[Option<SsTable>; NUM_CFS]) -> u64 {
        let mut total_blob_size = 0;
        for i in 0..NUM_CFS {
            if let Some(cf_tbl) = &cfs[i] {
                total_blob_size += cf_tbl.total_blob_size();
            }
        }
        total_blob_size
    }

    pub fn id(&self) -> u64 {
        self.file.id()
    }

    pub fn get_cf(&self, cf: usize) -> &Option<sstable::SsTable> {
        &self.cfs[cf]
    }

    pub fn size(&self) -> u64 {
        self.file.size()
    }

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

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

    pub fn tombs(&self) -> u64 {
        self.cfs[WRITE_CF].as_ref().map_or(0, |t| t.tombs as u64)
    }

    pub fn entries_write_cf(&self) -> u64 {
        self.cfs[WRITE_CF].as_ref().map_or(0, |t| t.entries as u64)
    }

    pub fn old_entries_write_cf(&self) -> u64 {
        self.cfs[WRITE_CF]
            .as_ref()
            .map_or(0, |t| t.old_entries as u64)
    }

    pub fn old_data_size_write_cf(&self) -> u64 {
        self.cfs[WRITE_CF].as_ref().map_or(0, |t| t.old_data_size())
    }

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

    pub fn smallest(&self) -> InnerKey<'_> {
        InnerKey::from_inner_buf(self.smallest.chunk())
    }

    pub fn biggest(&self) -> InnerKey<'_> {
        InnerKey::from_inner_buf(self.biggest.chunk())
    }

    pub fn snap_version(&self) -> SnapVersion {
        self.footer.snap_version
    }

    pub fn has_data_in_bound(&self, bound: DataBound<'_>) -> bool {
        if !self.data_bound().overlap_bound(bound) {
            return false;
        }
        self.cfs
            .iter()
            .filter_map(|t| t.as_ref())
            .any(|t| t.has_overlap(bound))
    }

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

    pub fn is_write_cf_only(&self) -> bool {
        self.footer.magic == MAGIC_NUMBER_SPLIT_L0
    }

    pub fn to_l0_create(&self) -> L0Create {
        let mut l0_create = L0Create::new();
        l0_create.set_id(self.id());
        l0_create.set_smallest(self.smallest().to_vec());
        l0_create.set_biggest(self.biggest().to_vec());
        l0_create.set_size(self.size() as u32);
        l0_create
    }
}

impl BoundedDataSet for L0TableCore {
    fn data_bound(&self) -> DataBound<'_> {
        DataBound::new(self.smallest(), self.biggest(), true)
    }
}

// Used to build L0Table from key value pairs.
//
// For detailed table format specification, refer to the comment above
// `L0Table`.
pub struct L0Builder {
    // A vector of builders used to construct the L0 table.
    // Each builder is responsible for adding key-value pairs and managing their respective states.
    builders: Vec<Builder>,
    version: SnapVersion,
    count: usize,
    fid: u64,
}

impl L0Builder {
    pub fn new(
        fid: u64,
        block_size: usize,
        version: SnapVersion,
        checksum_type: ChecksumType,
        encryption_key: Option<EncryptionKey>,
    ) -> Self {
        let mut builders = Vec::with_capacity(4);
        for _ in 0..NUM_CFS {
            let builder = Builder::new(
                fid,
                block_size,
                NO_COMPRESSION,
                0,
                checksum_type,
                encryption_key.clone(),
            );
            builders.push(builder);
        }
        Self {
            builders,
            version,
            count: 0,
            fid,
        }
    }

    pub fn add(
        &mut self,
        cf: usize,
        key: InnerKey<'_>,
        val: &Value,
        external_link: Option<BlobRef>,
    ) -> Result<()> {
        self.builders[cf].add(key, val, external_link)?;
        self.count += 1;
        Ok(())
    }

    pub fn finish(&mut self) -> Result<(L0Create, Bytes)> {
        let mut estimated_size = 0;
        for builder in &self.builders {
            estimated_size += builder.estimated_size();
        }
        let mut buf = Vec::with_capacity(estimated_size);
        let mut offsets = Vec::with_capacity(NUM_CFS);
        for builder in &mut self.builders {
            let offset = buf.len() as u32;
            offsets.push(offset);
            if !builder.is_empty() {
                builder.finish(offset, &mut buf)?;
            }
        }
        for offset in offsets {
            buf.put_u32_le(offset);
        }
        buf.put_u64_le(self.version.into_inner());
        buf.put_u32_le(NUM_CFS as u32);
        buf.put_u32_le(MAGIC_NUMBER);
        let (smallest, biggest) = self.smallest_biggest();
        let l0_create = new_l0_create_pb(self.fid, smallest, biggest, buf.len() as u32);
        Ok((l0_create, buf.into()))
    }

    fn smallest_biggest(&self) -> (Vec<u8>, Vec<u8>) {
        let mut smallest_buf = vec![];
        let mut biggest_buf = vec![];
        for builder in &self.builders {
            if !builder.get_smallest().is_empty()
                && (smallest_buf.is_empty() || builder.get_smallest() < smallest_buf.as_slice())
            {
                smallest_buf.truncate(0);
                smallest_buf.extend_from_slice(builder.get_smallest());
            }
            if builder.get_biggest() > biggest_buf.as_slice() {
                biggest_buf.truncate(0);
                biggest_buf.extend_from_slice(builder.get_biggest());
            }
        }
        (smallest_buf, biggest_buf)
    }

    pub fn total_blob_size(&self) -> u64 {
        let mut total_blob_size = 0;
        for builder in &self.builders {
            total_blob_size += builder.get_total_blob_size();
        }
        total_blob_size
    }

    pub fn is_empty(&self) -> bool {
        self.count == 0
    }

    pub fn get_fid(&self) -> u64 {
        self.fid
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::table::file::InMemFile;

    fn build_multi_cf_l0() -> Vec<u8> {
        let mut builder = L0Builder::new(1, 64, 1.into(), ChecksumType::Crc32c, None);
        let value = Value::new_with_version(b"value", 1);
        builder
            .add(
                WRITE_CF,
                InnerKey::from_inner_buf(b"write-key"),
                &value,
                None,
            )
            .unwrap();
        builder
            .add(
                crate::LOCK_CF,
                InnerKey::from_inner_buf(b"lock-key"),
                &value,
                None,
            )
            .unwrap();
        builder.finish().unwrap().1.to_vec()
    }

    fn cf_offsets_offset(data: &[u8]) -> usize {
        data.len() - L0_FOOTER_SIZE - L0_CF_OFFSETS_SIZE
    }

    fn assert_invalid_l0(data: Vec<u8>) {
        let file = Arc::new(InMemFile::new(1, data.into()));
        assert!(matches!(
            L0Table::new(file, BlockCache::None, false, EncryptionState::Disabled),
            Err(Error::InvalidFileSize)
        ));
    }

    #[test]
    fn test_l0_builder_basic() {
        // Test the creation of a new L0Builder instance
        let mut builder = L0Builder::new(1, 1024, 1.into(), ChecksumType::Crc32c, None);
        assert_eq!(builder.get_fid(), 1); // Ensure the fid is set correctly
        assert!(builder.is_empty()); // Ensure the builder is empty upon creation

        // Test adding a key-value pair to the L0Builder
        let key = InnerKey::from_inner_buf(b"test_key");
        let value = Value::new_with_version(b"test_value", 1);
        builder.add(0, key, &value, None).unwrap();
        assert!(!builder.is_empty()); // Ensure the builder is not empty after adding an entry

        // Test finishing the L0Builder and creating an L0Table
        let (l0_create, bytes) = builder.finish().unwrap();

        // Ensure the created L0Table has the expected properties
        assert_eq!(l0_create.get_id(), 1); // Check the ID of the created L0Table
        assert!(!bytes.is_empty()); // Ensure the byte array is not empty

        // Test is_empty on a new builder
        let mut builder2 = L0Builder::new(2, 1024, 1.into(), ChecksumType::Crc32c, None);
        assert!(builder2.is_empty()); // Ensure new builder is empty

        // Test is_empty after adding an entry
        let key = InnerKey::from_inner_buf(b"test_key2");
        let value = Value::new_with_version(b"test_value2", 1);
        builder2.add(0, key, &value, None).unwrap();
        assert!(!builder2.is_empty()); // Ensure the builder is not empty after adding an entry
    }

    #[test]
    fn test_l0_rejects_footer_outside_cf_segment() {
        let mut data = build_multi_cf_l0();

        let cf_offsets_offset = cf_offsets_offset(&data);
        let write_start = LittleEndian::read_u32(&data[cf_offsets_offset..]) as usize;
        let write_end = LittleEndian::read_u32(&data[cf_offsets_offset + 4..]) as usize;
        let write_size = write_end - write_start;
        let footer_offset = write_end - FOOTER_SIZE;
        let mut footer = Footer::default();
        footer.unmarshal(&data[footer_offset..write_end]);

        // Point the WRITE footer into the following CF. It is inside the L0 file but
        // outside the WRITE SST segment.
        footer.properties_offset = (write_size + 1) as u32;
        assert!(write_start + (footer.properties_offset as usize) < cf_offsets_offset);
        let mut footer_data = Vec::with_capacity(FOOTER_SIZE);
        footer.marshal(&mut footer_data);
        data[footer_offset..write_end].copy_from_slice(&footer_data);

        let file = Arc::new(InMemFile::new(1, data.into()));
        assert!(matches!(
            L0Table::new(file, BlockCache::None, false, EncryptionState::Disabled),
            Err(Error::InvalidFileSize)
        ));
    }

    #[test]
    fn test_l0_rejects_truncated_outer_metadata() {
        assert_invalid_l0(vec![0; L0_FOOTER_SIZE - 1]);

        let mut footer_only = Vec::with_capacity(L0_FOOTER_SIZE);
        footer_only.put_u64_le(1);
        footer_only.put_u32_le(NUM_CFS as u32);
        footer_only.put_u32_le(MAGIC_NUMBER);
        assert_invalid_l0(footer_only);
    }

    #[test]
    fn test_l0_rejects_wrong_outer_cf_count() {
        let mut data = build_multi_cf_l0();
        let footer_offset = data.len() - L0_FOOTER_SIZE;
        LittleEndian::write_u32(
            &mut data[footer_offset + 8..footer_offset + 12],
            NUM_CFS as u32 - 1,
        );
        assert_invalid_l0(data);
    }

    #[test]
    fn test_l0_rejects_invalid_outer_cf_offsets() {
        let mut reversed = build_multi_cf_l0();
        let offsets_offset = cf_offsets_offset(&reversed);
        let second = LittleEndian::read_u32(&reversed[offsets_offset + 4..]);
        let third = LittleEndian::read_u32(&reversed[offsets_offset + 8..]);
        assert!(second < third);
        LittleEndian::write_u32(&mut reversed[offsets_offset + 4..], third);
        LittleEndian::write_u32(&mut reversed[offsets_offset + 8..], second);
        assert_invalid_l0(reversed);

        let mut out_of_bounds = build_multi_cf_l0();
        let offsets_offset = cf_offsets_offset(&out_of_bounds);
        LittleEndian::write_u32(
            &mut out_of_bounds[offsets_offset + 8..],
            (offsets_offset + 1) as u32,
        );
        assert_invalid_l0(out_of_bounds);

        let mut nonzero_first = build_multi_cf_l0();
        let offsets_offset = cf_offsets_offset(&nonzero_first);
        let second = LittleEndian::read_u32(&nonzero_first[offsets_offset + 4..]);
        assert!(second > 0);
        LittleEndian::write_u32(&mut nonzero_first[offsets_offset..], second);
        assert_invalid_l0(nonzero_first);

        let mut all_empty = build_multi_cf_l0();
        let offsets_offset = cf_offsets_offset(&all_empty);
        for i in 0..NUM_CFS {
            LittleEndian::write_u32(
                &mut all_empty[offsets_offset + i * 4..],
                offsets_offset as u32,
            );
        }
        assert_invalid_l0(all_empty);
    }

    #[test]
    fn test_l0_cf_rejects_invalid_segment_bounds() {
        let file = Arc::new(InMemFile::new(1, build_multi_cf_l0().into()));
        assert!(matches!(
            SsTable::new_l0_cf(
                file.clone(),
                1,
                0,
                BlockCache::None,
                EncryptionState::Disabled,
            ),
            Err(Error::InvalidFileSize)
        ));
        assert!(matches!(
            SsTable::new_l0_cf(
                file.clone(),
                0,
                file.size() + 1,
                BlockCache::None,
                EncryptionState::Disabled,
            ),
            Err(Error::InvalidFileSize)
        ));
    }

    #[test]
    fn test_l0_footer_default() {
        // Test the default values of L0Footer.
        let footer = L0Footer::default();
        assert_eq!(footer.snap_version.into_inner(), 0); // Ensure the default version is 0.
        assert_eq!(footer.num_cfs, 0); // Ensure the default number of column families is 0.
        assert_eq!(footer.magic, 0); // Ensure the default magic number is 0.
    }

    #[test]
    fn test_l0_footer_unmarshal() {
        // Test the unmarshal method of L0Footer.
        let mut footer = L0Footer::default();
        let data: [u8; 16] = [
            1, 0, 0, 0, 0, 0, 0, 0, // version = 1
            2, 0, 0, 0, // num_cfs = 2
            0xEF, 0xBE, 0xAD, 0xDE, // magic = 0xDEADBEEF in little endian
        ];
        footer.unmarshal(&data);

        assert_eq!(footer.snap_version.into_inner(), 1); // Check that version is set correctly.
        assert_eq!(footer.num_cfs, 2); // Check that num_cfs is set correctly.
        assert_eq!(footer.magic, 0xDEADBEEF); // Check that magic number is set correctly.
    }

    #[test]
    fn test_l0_footer_is_match() {
        // Test the is_match method of L0Footer.
        let mut footer = L0Footer::default();

        // Test with a valid magic number.
        footer.magic = MAGIC_NUMBER;
        assert!(footer.is_match()); // Ensure it matches with MAGIC_NUMBER.

        // Test with another valid magic number.
        footer.magic = MAGIC_NUMBER_SPLIT_L0;
        assert!(footer.is_match()); // Ensure it matches with MAGIC_NUMBER_SPLIT_L0.

        // Test with an invalid magic number.
        footer.magic = 0x12345678;
        assert!(!footer.is_match()); // Ensure it does not match with an invalid magic number.
    }
}