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

use std::{
    cmp,
    iter::Iterator as StdIterator,
    ops::Deref,
    sync::{
        Arc, Mutex,
        atomic::{AtomicBool, AtomicUsize, Ordering},
    },
};

use super::{Arena, MemTableImplementation, SkipList};
use crate::{
    EXTRA_CF, NUM_CFS, WRITE_CF,
    table::{AtomicSnapVersion, DataBound, SnapVersion, TxnFile, memtable::skl_ext::SkipListExt},
};

/// One memtable generation across all column families.
///
/// Arena-backed column families share one arena. Crossbeam-backed column
/// families own independently reclaimable maps. New generations, splits, and
/// derived tables must preserve the selected implementation.
#[derive(Clone)]
pub struct CfTable {
    pub core: Arc<CfTableCore>,
}

impl Deref for CfTable {
    type Target = CfTableCore;

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

impl Default for CfTable {
    fn default() -> Self {
        Self::new()
    }
}

impl CfTable {
    pub fn new() -> Self {
        Self::with_implementation(MemTableImplementation::default())
    }

    pub(crate) fn with_implementation(implementation: MemTableImplementation) -> Self {
        Self {
            core: Arc::new(CfTableCore::with_implementation(implementation)),
        }
    }

    /// Creates a table with the requested backend for tests and benchmarks.
    #[cfg(feature = "testexport")]
    #[doc(hidden)]
    pub fn with_crossbeam_memtable_for_test(enabled: bool) -> Self {
        Self::with_implementation(MemTableImplementation::from_crossbeam_enabled(enabled))
    }

    /// Creates an empty successor with the same memtable implementation.
    pub(crate) fn new_empty_like(&self) -> Self {
        Self::with_implementation(self.core.implementation)
    }

    pub(crate) fn implementation(&self) -> MemTableImplementation {
        self.core.implementation
    }

    /// Permanently marks every column family in this generation immutable.
    pub(crate) fn seal(&self) {
        for table in &self.core.tbls {
            table.seal();
        }
    }

    /// Exposes the production sealing transition to benchmarks.
    #[cfg(feature = "testexport")]
    #[doc(hidden)]
    pub fn seal_for_test(&self) {
        self.seal();
    }

    pub fn new_split(&self) -> Self {
        let tbls = self.core.tbls.clone();
        let snap_ver = AtomicSnapVersion::new(self.snap_ver.load());
        let force_switch = AtomicBool::new(self.force_switch.load(Ordering::Acquire));
        let props = Mutex::new(self.core.props.lock().unwrap().clone());
        let unpersisted_props_size = AtomicUsize::new(self.core.unpersisted_props_size());
        Self {
            core: Arc::new(CfTableCore {
                tbls,
                implementation: self.core.implementation,
                snap_ver,
                force_switch,
                props,
                unpersisted_props_size,
            }),
        }
    }

    #[must_use]
    pub fn add_write_cf_txn_files(&self, txn_files: &[TxnFile]) -> Self {
        let mut tbls = self.core.tbls.clone();
        tbls[WRITE_CF] = tbls[WRITE_CF].add_txn_files(txn_files);
        let snap_ver = AtomicSnapVersion::new(self.snap_ver.load());
        let force_switch = AtomicBool::new(self.force_switch.load(Ordering::Acquire));
        let props = Mutex::new(self.core.props.lock().unwrap().clone());
        let unpersisted_props_size = AtomicUsize::new(self.core.unpersisted_props_size());
        Self {
            core: Arc::new(CfTableCore {
                tbls,
                implementation: self.core.implementation,
                snap_ver,
                force_switch,
                props,
                unpersisted_props_size,
            }),
        }
    }
}

pub struct CfTableCore {
    tbls: [SkipListExt; NUM_CFS],
    implementation: MemTableImplementation,
    snap_ver: AtomicSnapVersion,
    force_switch: AtomicBool,
    props: Mutex<Option<kvenginepb::Properties>>,
    unpersisted_props_size: AtomicUsize,
}

impl Default for CfTableCore {
    fn default() -> Self {
        Self::new()
    }
}

impl CfTableCore {
    pub fn new() -> Self {
        Self::with_implementation(MemTableImplementation::default())
    }

    fn with_implementation(implementation: MemTableImplementation) -> Self {
        let arena =
            (implementation == MemTableImplementation::Arena).then(|| Arc::new(Arena::new()));
        Self {
            tbls: [
                SkipListExt::new(SkipList::with_implementation(implementation, arena.clone())),
                SkipListExt::new(SkipList::with_implementation(implementation, arena.clone())),
                SkipListExt::new(SkipList::with_implementation(implementation, arena)),
            ],
            implementation,
            snap_ver: AtomicSnapVersion::new(SnapVersion::zero()),
            force_switch: AtomicBool::new(false),
            props: Mutex::new(None),
            unpersisted_props_size: AtomicUsize::default(),
        }
    }

    pub fn get_cf(&self, cf: usize) -> &SkipListExt {
        &self.tbls[cf]
    }

    pub fn is_empty(&self) -> bool {
        for tbl in &self.tbls {
            if !tbl.is_empty() {
                return false;
            }
        }
        true
    }

    /// Returns deterministic cumulative submitted bytes across all CFs.
    pub fn size(&self) -> u64 {
        self.tbls.iter().map(|t| t.size() as u64).sum()
    }

    /// Returns retained-memory charges across all column families.
    pub fn memory_usage(&self) -> u64 {
        self.tbls.iter().map(|t| t.memory_usage() as u64).sum()
    }

    pub fn skip_list_entries(&self) -> usize {
        self.tbls.iter().map(|t| t.skip_list_entries()).sum()
    }

    pub fn skip_list_size(&self) -> usize {
        self.tbls.iter().map(|t| t.skl_size()).sum()
    }

    pub fn set_snap_version(&self, snap_ver: SnapVersion) {
        self.snap_ver.store(snap_ver)
    }

    pub fn set_properties(&self, props: kvenginepb::Properties) {
        self.props.lock().unwrap().replace(props);
    }

    pub fn get_properties(&self) -> Option<kvenginepb::Properties> {
        self.props.lock().unwrap().clone()
    }

    pub fn add_unpersisted_props_size(&self, size: usize) {
        self.unpersisted_props_size
            .fetch_add(size, Ordering::Release);
    }

    pub fn unpersisted_props_size(&self) -> usize {
        self.unpersisted_props_size.load(Ordering::Acquire)
    }

    pub fn get_snap_version(&self) -> SnapVersion {
        self.snap_ver.load()
    }

    pub fn set_force_switch(&self) {
        self.force_switch.store(true, Ordering::Release);
    }

    pub fn is_force_switch(&self) -> bool {
        self.force_switch.load(Ordering::Acquire)
    }

    pub fn has_data_in_bound(&self, bound: DataBound<'_>) -> bool {
        if self.is_empty() {
            return false;
        }
        for cf in 0..NUM_CFS {
            let tbl = &self.tbls[cf];
            let mut iter = tbl.new_iterator(false);
            iter.seek(bound.lower_bound);
            if iter.valid() && !bound.less_than_key(iter.key()) {
                return true;
            }
        }
        false
    }

    pub fn data_max_ts(&self) -> u64 {
        // Ignore LOCK_CF, as `ts` in LOCK_CF is not a TSO.
        // TODO: in async_commit, LOCK_CF may contains data, need handle it later.
        cmp::max(
            self.tbls[WRITE_CF].data_max_ts(),
            self.tbls[EXTRA_CF].data_max_ts(),
        )
    }
}

#[cfg(feature = "debug-trace-mem-table")]
impl CfTableCore {
    pub fn skl_and_txn_files_size(&self) -> (u64 /* skl_size */, u64 /* txn_files_size */) {
        self.tbls
            .iter()
            .fold((0, 0), |(skl_size, txn_files_size), t| {
                (
                    skl_size + t.skl_size() as u64,
                    txn_files_size + t.txn_files_size() as u64,
                )
            })
    }
}

#[cfg(test)]
mod tests {
    use std::panic::{AssertUnwindSafe, catch_unwind};

    use rstest::rstest;

    use super::*;
    use crate::table::{InnerKey, memtable::WriteBatch};

    #[test]
    fn test_memtable_implementation_selection_and_split() {
        let crossbeam = CfTable::new();
        assert_eq!(
            crossbeam.implementation(),
            MemTableImplementation::Crossbeam
        );
        assert_eq!(
            crossbeam.new_empty_like().implementation(),
            MemTableImplementation::Crossbeam
        );
        assert_eq!(
            crossbeam.new_split().implementation(),
            MemTableImplementation::Crossbeam
        );

        let arena = CfTable::with_implementation(MemTableImplementation::Arena);
        assert_eq!(arena.implementation(), MemTableImplementation::Arena);
        assert_eq!(
            arena.new_empty_like().implementation(),
            MemTableImplementation::Arena
        );
        assert_eq!(
            arena.new_split().implementation(),
            MemTableImplementation::Arena
        );
        assert!(!SkipList::default().supports_get_hint());
    }

    /// Test basic operations on an empty CfTable
    #[rstest]
    #[case::arena(MemTableImplementation::Arena)]
    #[case::crossbeam(MemTableImplementation::Crossbeam)]
    fn test_empty_cftable(#[case] implementation: MemTableImplementation) {
        let table = CfTable::with_implementation(implementation);

        // Test empty state
        assert!(table.is_empty());
        assert_eq!(table.size(), 0);
        assert_eq!(table.skip_list_size(), 0);
        assert_eq!(table.get_snap_version().into_inner(), 0);
        assert!(!table.is_force_switch());

        // Test empty properties
        assert!(table.get_properties().is_none());
    }

    #[test]
    fn test_crossbeam_seal_is_idempotent_and_covers_all_cfs() {
        let table = CfTable::with_implementation(MemTableImplementation::Crossbeam);
        table.seal();
        table.seal();

        for cf in 0..NUM_CFS {
            let mut wb = WriteBatch::new();
            wb.put(InnerKey::from_inner_buf(b"key"), 0, &[], 1, b"value");
            assert!(
                catch_unwind(AssertUnwindSafe(|| table
                    .get_cf(cf)
                    .put_batch(&mut wb, None, cf)))
                .is_err()
            );
        }
    }

    #[test]
    fn test_arena_seal_preserves_iterator_and_write_behavior() {
        let table = CfTable::with_implementation(MemTableImplementation::Arena);
        let mut wb = WriteBatch::new();
        wb.put(InnerKey::from_inner_buf(b"key"), 0, &[], 1, b"value-1");
        table.get_cf(WRITE_CF).put_batch(&mut wb, None, WRITE_CF);
        table.seal();

        let mut regular = table.get_cf(WRITE_CF).new_iterator(false);
        let mut flush = table.get_cf(WRITE_CF).new_flush_iterator();
        regular.rewind();
        flush.rewind();
        assert_eq!(regular.key(), flush.key());
        assert_eq!(regular.value().get_value(), flush.value().get_value());

        let mut wb = WriteBatch::new();
        wb.put(InnerKey::from_inner_buf(b"key"), 0, &[], 2, b"value-2");
        table.get_cf(WRITE_CF).put_batch(&mut wb, None, WRITE_CF);
        let mut iterator = table.get_cf(WRITE_CF).new_flush_iterator();
        iterator.rewind();
        assert_eq!(iterator.value().get_value(), b"value-2");
    }

    /// Test version management
    #[rstest]
    #[case::arena(MemTableImplementation::Arena)]
    #[case::crossbeam(MemTableImplementation::Crossbeam)]
    fn test_version_management(#[case] implementation: MemTableImplementation) {
        let table = CfTable::with_implementation(implementation);

        // Set and verify version
        table.set_snap_version(SnapVersion::from(42));
        assert_eq!(table.get_snap_version().into_inner(), 42);

        // Create split table and verify version inheritance
        let split_table = table.new_split();
        assert_eq!(split_table.get_snap_version().into_inner(), 42);

        // Verify version independence after split
        split_table.set_snap_version(SnapVersion::from(100));
        assert_eq!(split_table.get_snap_version().into_inner(), 100);
        assert_eq!(table.get_snap_version().into_inner(), 42);
    }

    /// Test force switch functionality
    #[rstest]
    #[case::arena(MemTableImplementation::Arena)]
    #[case::crossbeam(MemTableImplementation::Crossbeam)]
    fn test_force_switch(#[case] implementation: MemTableImplementation) {
        let table = CfTable::with_implementation(implementation);
        assert!(!table.is_force_switch());

        table.set_force_switch();
        assert!(table.is_force_switch());

        // Test force switch in split table
        let split_table = table.new_split();
        assert!(split_table.is_force_switch());
    }

    /// Test data operations across different CFs
    #[rstest]
    #[case::arena(MemTableImplementation::Arena)]
    #[case::crossbeam(MemTableImplementation::Crossbeam)]
    fn test_cf_operations(#[case] implementation: MemTableImplementation) {
        let table = CfTable::with_implementation(implementation);

        // Write data to different CFs
        let mut kv_size = 0;
        for cf in 0..NUM_CFS {
            let tbl = table.get_cf(cf);
            let mut wb = WriteBatch::new();
            let key = format!("key_{}", cf);
            let value = format!("value_{}", cf);
            kv_size += key.len() + value.len();
            wb.put(
                InnerKey::from_inner_buf(key.as_bytes()),
                0,
                &[0],
                1,
                value.as_bytes(),
            );
            tbl.put_batch(&mut wb, None, cf);
        }

        // Verify data presence
        assert!(!table.is_empty());
        assert!(table.size() > 0);
        assert!(
            table.skip_list_size() >= kv_size,
            "{:?}",
            (table.skip_list_size(), kv_size)
        );

        // Check data in each CF
        for cf in 0..NUM_CFS {
            let tbl = table.get_cf(cf);
            let mut it = tbl.new_iterator(false);
            it.rewind();
            assert!(it.valid());
            assert_eq!(it.key().deref(), format!("key_{}", cf).as_bytes());
        }
    }

    /// Test data bound checking
    #[rstest]
    #[case::arena(MemTableImplementation::Arena)]
    #[case::crossbeam(MemTableImplementation::Crossbeam)]
    fn test_data_bound_checking(#[case] implementation: MemTableImplementation) {
        let table = CfTable::with_implementation(implementation);

        // Insert test data
        let cf = WRITE_CF;
        let tbl = table.get_cf(cf);
        let mut wb = WriteBatch::new();
        wb.put(InnerKey::from_inner_buf(b"key_50"), 0, &[0], 1, b"value_50");
        tbl.put_batch(&mut wb, None, cf);

        // Test bounds
        let bound = DataBound {
            lower_bound: InnerKey::from_inner_buf(b"key_00"),
            upper_bound: InnerKey::from_inner_buf(b"key_99"),
            upper_inclusive: true,
        };
        assert!(table.has_data_in_bound(bound));

        let bound = DataBound {
            lower_bound: InnerKey::from_inner_buf(b"key_60"),
            upper_bound: InnerKey::from_inner_buf(b"key_99"),
            upper_inclusive: true,
        };
        assert!(!table.has_data_in_bound(bound));
    }

    /// Test properties management
    #[rstest]
    #[case::arena(MemTableImplementation::Arena)]
    #[case::crossbeam(MemTableImplementation::Crossbeam)]
    fn test_properties(#[case] implementation: MemTableImplementation) {
        let table = CfTable::with_implementation(implementation);

        // Initially no properties
        assert!(table.get_properties().is_none());

        // Set and get properties
        let props = kvenginepb::Properties::default();
        table.set_properties(props.clone());

        // Test properties in split table
        let split_table = table.new_split();
        assert!(split_table.get_properties().is_some());
    }

    /// Test data max timestamp
    #[rstest]
    #[case::arena(MemTableImplementation::Arena)]
    #[case::crossbeam(MemTableImplementation::Crossbeam)]
    fn test_data_max_ts(#[case] implementation: MemTableImplementation) {
        let table = CfTable::with_implementation(implementation);

        // Write data with different timestamps to different CFs
        let write_cf = table.get_cf(WRITE_CF);
        let mut wb = WriteBatch::new();
        wb.put(
            InnerKey::from_inner_buf(b"key1"),
            0,
            &[0],
            100, // timestamp
            b"value1",
        );
        write_cf.put_batch(&mut wb, None, WRITE_CF);

        let extra_cf = table.get_cf(EXTRA_CF);
        wb.reset();
        wb.put(
            InnerKey::from_inner_buf(b"key2"),
            0,
            &[0],
            200, // timestamp
            b"value2",
        );
        extra_cf.put_batch(&mut wb, None, EXTRA_CF);

        // Max timestamp should be 200
        assert_eq!(table.data_max_ts(), 200);
    }
}