use lsm_tree::compaction::PullDown;
use lsm_tree::compaction::filter::{
    CompactionFilter, Context as CompactionFilterContext, Factory, ItemAccessor, Verdict,
};
use lsm_tree::{AbstractTree, SeqNo, SequenceNumberCounter, get_tmp_folder};
use std::sync::{Arc, Mutex};
use test_log::test;

fn u32_s(n: u32) -> [u8; 4] {
    n.to_be_bytes()
}

fn u32_f(mut buf: &[u8]) -> u32 {
    use byteorder::{BE, ReadBytesExt};

    buf.read_u32::<BE>().unwrap()
}

fn filter_basic() -> lsm_tree::Result<()> {
    struct FilterState {
        disable: bool,
        do_rewrite: bool,
        saw_value: bool,
        drop_threshold: u32,
        finished: bool,
        expect_last_level: Option<bool>,
    }

    let filter_state = Arc::new(Mutex::new(FilterState {
        disable: true,
        do_rewrite: false,
        saw_value: false,
        drop_threshold: 4,
        finished: false,
        expect_last_level: None,
    }));

    struct Filter(Arc<Mutex<FilterState>>);

    impl CompactionFilter for Filter {
        fn filter_item(
            &mut self,
            item: ItemAccessor<'_>,
            _ctx: &CompactionFilterContext,
        ) -> lsm_tree::Result<Verdict> {
            let mut state = self.0.lock().unwrap();
            if state.disable {
                return Ok(Verdict::Keep);
            }

            let key = u32_f(item.key());

            if key == 0x00abcdef {
                Ok(Verdict::RemoveWeak)
            } else if key == 0x01abcdef {
                Ok(Verdict::Destroy)
            } else if key >= 0xff000000 {
                let value = u32_f(&item.value()?);
                assert_eq!(key & 0xff, value);
                state.saw_value = true;

                if !state.do_rewrite {
                    Ok(Verdict::Keep)
                } else {
                    Ok(Verdict::ReplaceValue(vec![b'a'; value as usize].into()))
                }
            } else {
                assert_eq!(item.value().expect("failed fetch"), b"a");
                if key < state.drop_threshold {
                    Ok(Verdict::Remove)
                } else {
                    Ok(Verdict::Keep)
                }
            }
        }

        fn finish(self: Box<Self>) {
            let mut state = self.0.lock().unwrap();
            state.finished = true;
        }
    }

    struct FilterFactory(Arc<Mutex<FilterState>>);

    impl Factory for FilterFactory {
        fn name(&self) -> &str {
            "Test"
        }

        fn make_filter(&self, ctx: &CompactionFilterContext) -> Box<dyn CompactionFilter> {
            {
                let guard = self.0.lock().unwrap();
                if let Some(expect_last_level) = guard.expect_last_level {
                    assert_eq!(ctx.is_last_level, expect_last_level)
                }
            }
            Box::new(Filter(self.0.clone()))
        }
    }

    let folder = get_tmp_folder();

    let mut config = lsm_tree::Config::new(
        &folder,
        SequenceNumberCounter::default(),
        SequenceNumberCounter::default(),
    )
    .with_compaction_filter_factory(Some(Arc::new(FilterFactory(filter_state.clone()))));
    config.level_count = 3;
    let tree = config.open()?;

    for i in 0u32..255 {
        tree.insert(u32_s(i), "old", 0);
    }

    tree.flush_active_memtable(0)?;
    tree.major_compact(u64::MAX, SeqNo::MAX)?;

    for i in 0u32..255 {
        tree.insert(u32_s(i), "a", 0);
        tree.insert(u32_s(i | 0xff000000), u32_s(i), 0);
    }

    tree.insert(u32_s(0x00abcdef), "test", 0);
    tree.insert(u32_s(0x01abcdef), "test", 0);

    tree.flush_active_memtable(0)?;

    let mut state = filter_state.lock().unwrap();
    state.disable = false;
    state.expect_last_level = Some(false);
    drop(state);

    tree.compact(Arc::new(PullDown(0, 1)), SeqNo::MAX)?;

    // filter should have dropped this
    assert!(tree.get(u32_s(2), SeqNo::MAX)?.is_none());

    // but not these
    assert!(tree.get(u32_s(5), SeqNo::MAX)?.is_some());
    assert!(tree.get(u32_s(12), SeqNo::MAX)?.is_some());

    // ensure weak-deleted value is gone
    assert!(tree.get(u32_s(0x00abcdef), SeqNo::MAX)?.is_none());
    // ensure destroyed value is gone
    assert!(tree.get(u32_s(0x01abcdef), SeqNo::MAX)?.is_none());

    let mut state = filter_state.lock().unwrap();
    // filter should think it was called
    assert!(state.saw_value);
    assert!(state.finished);
    state.saw_value = false;
    // up the threshold
    state.drop_threshold = 8;
    state.do_rewrite = true;
    state.expect_last_level = Some(true);
    drop(state);

    // compact to last level
    tree.major_compact(u64::MAX, SeqNo::MAX)?;

    // verify values gone
    assert!(tree.get(u32_s(2), SeqNo::MAX)?.is_none());
    assert!(tree.get(u32_s(5), SeqNo::MAX)?.is_none());
    // ensure the other value isn't
    assert!(tree.get(u32_s(12), SeqNo::MAX)?.is_some());

    // verify rewrites work
    assert_eq!(tree.get(u32_s(2 | 0xff000000), SeqNo::MAX)?.unwrap(), b"aa");
    assert_eq!(
        tree.get(u32_s(37 | 0xff000000), SeqNo::MAX)?.unwrap(),
        b"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
    );

    let state = filter_state.lock().unwrap();
    assert!(state.saw_value);
    drop(state);

    Ok(())
}

#[test]
fn compaction_filter() -> lsm_tree::Result<()> {
    filter_basic()
}
