#![cfg(feature = "proptest")]

#[cfg(feature = "var-collections")]
use armdb::VarTree;
use armdb::{Config, ConstTree, DbError};
use proptest::prelude::*;
use std::collections::BTreeMap;
use tempfile::TempDir;

// --- Helpers ---

fn make_key(group: u8, id: u16) -> [u8; 8] {
    let mut k = [0u8; 8];
    k[0] = group;
    k[6..8].copy_from_slice(&id.to_be_bytes());
    k
}

fn make_config(reversed: bool) -> Config {
    let mut config = Config::test();
    config.shard_count = 2;
    config.max_file_size = 4096;
    config.write_buffer_size = 4096;
    config.compaction_threshold = 0.1;
    config.reversed = reversed;
    config
}

// --- ConstTree ---

#[derive(Debug, Clone)]
enum ConstOp {
    Put {
        group: u8,
        id: u16,
        value: u64,
    },
    Get {
        group: u8,
        id: u16,
    },
    Delete {
        group: u8,
        id: u16,
    },
    Insert {
        group: u8,
        id: u16,
        value: u64,
    },
    Cas {
        group: u8,
        id: u16,
        expected: u64,
        new_value: u64,
    },
    Compact,
    CloseReopen,
    IterAll,
    PrefixIter {
        group: u8,
    },
}

fn const_op_strategy() -> impl Strategy<Value = ConstOp> {
    prop_oneof![
        40 => (0..4u8, 0..16u16, any::<u64>())
            .prop_map(|(g, id, v)| ConstOp::Put { group: g, id, value: v }),
        20 => (0..4u8, 0..16u16)
            .prop_map(|(g, id)| ConstOp::Get { group: g, id }),
        15 => (0..4u8, 0..16u16)
            .prop_map(|(g, id)| ConstOp::Delete { group: g, id }),
        5 => (0..4u8, 0..16u16, any::<u64>())
            .prop_map(|(g, id, v)| ConstOp::Insert { group: g, id, value: v }),
        5 => (0..4u8, 0..16u16, any::<u64>(), any::<u64>())
            .prop_map(|(g, id, e, n)| ConstOp::Cas { group: g, id, expected: e, new_value: n }),
        5 => Just(ConstOp::Compact),
        3 => Just(ConstOp::CloseReopen),
        5 => Just(ConstOp::IterAll),
        2 => (0..4u8).prop_map(|g| ConstOp::PrefixIter { group: g }),
    ]
}

fn run_const_tree_model(ops: Vec<ConstOp>, reversed: bool) {
    let dir = TempDir::new().unwrap();
    let mut model: BTreeMap<[u8; 8], [u8; 8]> = BTreeMap::new();
    let mut tree = ConstTree::<[u8; 8], 8>::open(dir.path(), make_config(reversed)).unwrap();

    for (i, op) in ops.iter().enumerate() {
        match op {
            ConstOp::Put { group, id, value } => {
                let k = make_key(*group, *id);
                let v = value.to_be_bytes();
                let tree_old = tree.put(&k, &v).unwrap();
                let model_old = model.insert(k, v);
                assert_eq!(tree_old, model_old, "put old mismatch at op {i}");
            }
            ConstOp::Get { group, id } => {
                let k = make_key(*group, *id);
                let tree_val = tree.get(&k);
                let model_val = model.get(&k).copied();
                assert_eq!(tree_val, model_val, "get mismatch at op {i}");
            }
            ConstOp::Delete { group, id } => {
                let k = make_key(*group, *id);
                let tree_old = tree.delete(&k).unwrap();
                let model_old = model.remove(&k);
                assert_eq!(tree_old, model_old, "delete old mismatch at op {i}");
            }
            ConstOp::Insert { group, id, value } => {
                let k = make_key(*group, *id);
                let v = value.to_be_bytes();
                let result = tree.insert(&k, &v);
                #[allow(clippy::map_entry)]
                if model.contains_key(&k) {
                    assert!(
                        matches!(result, Err(DbError::KeyExists)),
                        "insert on existing key should be KeyExists at op {i}, got {result:?}"
                    );
                } else {
                    result.unwrap();
                    model.insert(k, v);
                }
            }
            ConstOp::Cas {
                group,
                id,
                expected,
                new_value,
            } => {
                let k = make_key(*group, *id);
                let e = expected.to_be_bytes();
                let n = new_value.to_be_bytes();
                let result = tree.cas(&k, &e, &n);
                match model.get(&k) {
                    None => assert!(
                        matches!(result, Err(DbError::KeyNotFound)),
                        "cas on missing key should be KeyNotFound at op {i}, got {result:?}"
                    ),
                    Some(v) if *v == e => {
                        result.unwrap();
                        model.insert(k, n);
                    }
                    Some(_) => assert!(
                        matches!(result, Err(DbError::CasMismatch)),
                        "cas mismatch should be CasMismatch at op {i}, got {result:?}"
                    ),
                }
            }
            ConstOp::Compact => {
                tree.compact().unwrap();
                assert_eq!(
                    tree.len(),
                    model.len(),
                    "len mismatch after compact at op {i}"
                );
            }
            ConstOp::CloseReopen => {
                tree.close().unwrap();
                tree = ConstTree::<[u8; 8], 8>::open(dir.path(), make_config(reversed)).unwrap();
                assert_eq!(
                    tree.len(),
                    model.len(),
                    "len mismatch after reopen at op {i}"
                );
                for (k, v) in &model {
                    assert_eq!(
                        tree.get(k),
                        Some(*v),
                        "value mismatch after reopen at op {i}, key {k:?}"
                    );
                }
            }
            ConstOp::IterAll => {
                let tree_entries: Vec<_> = tree.iter().collect();
                let model_entries: Vec<_> = if reversed {
                    model.iter().rev().map(|(k, v)| (*k, *v)).collect()
                } else {
                    model.iter().map(|(k, v)| (*k, *v)).collect()
                };
                assert_eq!(tree_entries, model_entries, "iter mismatch at op {i}");
            }
            ConstOp::PrefixIter { group } => {
                let tree_entries: Vec<_> = tree.prefix_iter(&[*group]).collect();
                let model_entries: Vec<_> = if reversed {
                    model
                        .iter()
                        .rev()
                        .filter(|(k, _)| k[0] == *group)
                        .map(|(k, v)| (*k, *v))
                        .collect()
                } else {
                    model
                        .iter()
                        .filter(|(k, _)| k[0] == *group)
                        .map(|(k, v)| (*k, *v))
                        .collect()
                };
                assert_eq!(
                    tree_entries, model_entries,
                    "prefix_iter mismatch at op {i}"
                );
            }
        }
    }

    // Final consistency check
    assert_eq!(tree.len(), model.len(), "final len mismatch");
    for (k, v) in &model {
        assert_eq!(tree.get(k).unwrap(), *v, "final value mismatch for {k:?}");
    }
}

// --- VarTree ---

#[cfg(feature = "var-collections")]
#[derive(Debug, Clone)]
enum VarOp {
    Put {
        group: u8,
        id: u16,
        value: Vec<u8>,
    },
    Get {
        group: u8,
        id: u16,
    },
    Delete {
        group: u8,
        id: u16,
    },
    Insert {
        group: u8,
        id: u16,
        value: Vec<u8>,
    },
    Cas {
        group: u8,
        id: u16,
        expected: Vec<u8>,
        new_value: Vec<u8>,
    },
    Compact,
    CloseReopen,
    IterAll,
    PrefixIter {
        group: u8,
    },
}

#[cfg(feature = "var-collections")]
fn var_op_strategy() -> impl Strategy<Value = VarOp> {
    prop_oneof![
        40 => (0..4u8, 0..16u16, prop::collection::vec(any::<u8>(), 0..128))
            .prop_map(|(g, id, v)| VarOp::Put { group: g, id, value: v }),
        20 => (0..4u8, 0..16u16)
            .prop_map(|(g, id)| VarOp::Get { group: g, id }),
        15 => (0..4u8, 0..16u16)
            .prop_map(|(g, id)| VarOp::Delete { group: g, id }),
        5 => (0..4u8, 0..16u16, prop::collection::vec(any::<u8>(), 0..128))
            .prop_map(|(g, id, v)| VarOp::Insert { group: g, id, value: v }),
        5 => (0..4u8, 0..16u16, prop::collection::vec(any::<u8>(), 0..128), prop::collection::vec(any::<u8>(), 0..128))
            .prop_map(|(g, id, e, n)| VarOp::Cas { group: g, id, expected: e, new_value: n }),
        5 => Just(VarOp::Compact),
        3 => Just(VarOp::CloseReopen),
        5 => Just(VarOp::IterAll),
        2 => (0..4u8).prop_map(|g| VarOp::PrefixIter { group: g }),
    ]
}

#[cfg(feature = "var-collections")]
fn run_var_tree_model(ops: Vec<VarOp>, reversed: bool) {
    let dir = TempDir::new().unwrap();
    let mut model: BTreeMap<[u8; 8], Vec<u8>> = BTreeMap::new();
    let mut tree = VarTree::<[u8; 8]>::open(dir.path(), make_config(reversed)).unwrap();

    for (i, op) in ops.iter().enumerate() {
        match op {
            VarOp::Put { group, id, value } => {
                let k = make_key(*group, *id);
                tree.put(&k, value).unwrap();
                model.insert(k, value.clone());
            }
            VarOp::Get { group, id } => {
                let k = make_key(*group, *id);
                let tree_val = tree.get(&k).map(|bv| bv.as_ref().to_vec());
                let model_val = model.get(&k).cloned();
                assert_eq!(tree_val, model_val, "get mismatch at op {i}");
            }
            VarOp::Delete { group, id } => {
                let k = make_key(*group, *id);
                let tree_existed = tree.delete(&k).unwrap();
                let model_existed = model.remove(&k).is_some();
                assert_eq!(
                    tree_existed, model_existed,
                    "delete existed mismatch at op {i}"
                );
            }
            VarOp::Insert { group, id, value } => {
                let k = make_key(*group, *id);
                let result = tree.insert(&k, value);
                #[allow(clippy::map_entry)]
                if model.contains_key(&k) {
                    assert!(
                        matches!(result, Err(DbError::KeyExists)),
                        "insert on existing key should be KeyExists at op {i}, got {result:?}"
                    );
                } else {
                    result.unwrap();
                    model.insert(k, value.clone());
                }
            }
            VarOp::Cas {
                group,
                id,
                expected,
                new_value,
            } => {
                let k = make_key(*group, *id);
                let result = tree.cas(&k, expected, new_value);
                match model.get(&k) {
                    None => assert!(
                        matches!(result, Err(DbError::KeyNotFound)),
                        "cas on missing key should be KeyNotFound at op {i}, got {result:?}"
                    ),
                    Some(v) if v.as_slice() == expected.as_slice() => {
                        result.unwrap();
                        model.insert(k, new_value.clone());
                    }
                    Some(_) => assert!(
                        matches!(result, Err(DbError::CasMismatch)),
                        "cas mismatch should be CasMismatch at op {i}, got {result:?}"
                    ),
                }
            }
            VarOp::Compact => {
                tree.compact().unwrap();
                assert_eq!(
                    tree.len(),
                    model.len(),
                    "len mismatch after compact at op {i}"
                );
            }
            VarOp::CloseReopen => {
                tree.close().unwrap();
                tree = VarTree::open(dir.path(), make_config(reversed)).unwrap();
                assert_eq!(
                    tree.len(),
                    model.len(),
                    "len mismatch after reopen at op {i}"
                );
                for (k, v) in &model {
                    assert_eq!(
                        tree.get(k).unwrap().as_ref(),
                        v.as_slice(),
                        "value mismatch after reopen at op {i}, key {k:?}"
                    );
                }
            }
            VarOp::IterAll => {
                let tree_entries: Vec<_> = tree
                    .iter()
                    .map(|(k, bv)| (k, bv.as_ref().to_vec()))
                    .collect();
                let model_entries: Vec<_> = if reversed {
                    model.iter().rev().map(|(k, v)| (*k, v.clone())).collect()
                } else {
                    model.iter().map(|(k, v)| (*k, v.clone())).collect()
                };
                assert_eq!(tree_entries, model_entries, "iter mismatch at op {i}");
            }
            VarOp::PrefixIter { group } => {
                let tree_entries: Vec<_> = tree
                    .prefix_iter(&[*group])
                    .map(|(k, bv)| (k, bv.as_ref().to_vec()))
                    .collect();
                let model_entries: Vec<_> = if reversed {
                    model
                        .iter()
                        .rev()
                        .filter(|(k, _)| k[0] == *group)
                        .map(|(k, v)| (*k, v.clone()))
                        .collect()
                } else {
                    model
                        .iter()
                        .filter(|(k, _)| k[0] == *group)
                        .map(|(k, v)| (*k, v.clone()))
                        .collect()
                };
                assert_eq!(
                    tree_entries, model_entries,
                    "prefix_iter mismatch at op {i}"
                );
            }
        }
    }

    // Final consistency check
    assert_eq!(tree.len(), model.len(), "final len mismatch");
    for (k, v) in &model {
        assert_eq!(
            tree.get(k).unwrap().as_ref(),
            v.as_slice(),
            "final value mismatch for {k:?}"
        );
    }
}

// --- proptest ---

proptest! {
    #![proptest_config(ProptestConfig::with_cases(256))]

    #[test]
    fn const_tree_model_reversed(ops in prop::collection::vec(const_op_strategy(), 1..200)) {
        run_const_tree_model(ops, true);
    }

    #[test]
    fn const_tree_model_ascending(ops in prop::collection::vec(const_op_strategy(), 1..200)) {
        run_const_tree_model(ops, false);
    }
}

#[cfg(feature = "var-collections")]
proptest! {
    #![proptest_config(ProptestConfig::with_cases(256))]

    #[test]
    fn var_tree_model_reversed(ops in prop::collection::vec(var_op_strategy(), 1..200)) {
        run_var_tree_model(ops, true);
    }

    #[test]
    fn var_tree_model_ascending(ops in prop::collection::vec(var_op_strategy(), 1..200)) {
        run_var_tree_model(ops, false);
    }
}
