armdb 0.1.13

use armdb::{Config, MigrateAction, ZeroTree};
use tempfile::tempdir;
use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout};

// ---------------------------------------------------------------------------
// Test types
// ---------------------------------------------------------------------------

#[derive(
    FromBytes, IntoBytes, Immutable, KnownLayout, Clone, Copy, Debug, PartialEq, Eq, Default,
)]
#[repr(C)]
struct Counters {
    views: u64,
    likes: u64,
}

#[derive(
    FromBytes, IntoBytes, Immutable, KnownLayout, Clone, Copy, Debug, PartialEq, Eq, Default,
)]
#[repr(C)]
struct Point {
    x: i32,
    y: i32,
    z: i32,
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

type CounterTree = ZeroTree<[u8; 8], { size_of::<Counters>() }, Counters>;

fn open_counter_tree(dir: &std::path::Path) -> CounterTree {
    CounterTree::open(dir, Config::test()).unwrap()
}

fn populate(tree: &CounterTree, n: u64) {
    for i in 0..n {
        tree.put(
            &i.to_be_bytes(),
            &Counters {
                views: i * 10,
                likes: i,
            },
        )
        .unwrap();
    }
}

// ---------------------------------------------------------------------------
// Basic CRUD
// ---------------------------------------------------------------------------

#[test]
fn test_put_get() {
    let dir = tempdir().unwrap();
    let tree = open_counter_tree(dir.path());

    assert!(tree.get(&1u64.to_be_bytes()).is_none());

    tree.put(
        &1u64.to_be_bytes(),
        &Counters {
            views: 42,
            likes: 7,
        },
    )
    .unwrap();

    let c = tree.get(&1u64.to_be_bytes()).unwrap();
    assert_eq!(c.views, 42);
    assert_eq!(c.likes, 7);
}

#[test]
fn test_put_returns_old() {
    let dir = tempdir().unwrap();
    let tree = open_counter_tree(dir.path());

    let old = tree
        .put(
            &1u64.to_be_bytes(),
            &Counters {
                views: 10,
                likes: 1,
            },
        )
        .unwrap();
    assert!(old.is_none());

    let old = tree
        .put(
            &1u64.to_be_bytes(),
            &Counters {
                views: 20,
                likes: 2,
            },
        )
        .unwrap()
        .unwrap();
    assert_eq!(old.views, 10);
    assert_eq!(old.likes, 1);

    let current = tree.get(&1u64.to_be_bytes()).unwrap();
    assert_eq!(current.views, 20);
}

#[test]
fn test_insert() {
    let dir = tempdir().unwrap();
    let tree = open_counter_tree(dir.path());

    tree.insert(&1u64.to_be_bytes(), &Counters::default())
        .unwrap();

    let err = tree
        .insert(&1u64.to_be_bytes(), &Counters::default())
        .unwrap_err();
    assert!(matches!(err, armdb::DbError::KeyExists));
}

#[test]
fn test_delete() {
    let dir = tempdir().unwrap();
    let tree = open_counter_tree(dir.path());

    tree.put(
        &1u64.to_be_bytes(),
        &Counters {
            views: 99,
            likes: 5,
        },
    )
    .unwrap();

    let old = tree.delete(&1u64.to_be_bytes()).unwrap().unwrap();
    assert_eq!(old.views, 99);
    assert!(tree.get(&1u64.to_be_bytes()).is_none());
}

#[test]
fn test_contains_len() {
    let dir = tempdir().unwrap();
    let tree = open_counter_tree(dir.path());

    assert!(!tree.contains(&1u64.to_be_bytes()));
    assert_eq!(tree.len(), 0);
    assert!(tree.is_empty());

    populate(&tree, 5);
    assert!(tree.contains(&0u64.to_be_bytes()));
    assert_eq!(tree.len(), 5);
}

// ---------------------------------------------------------------------------
// CAS / Update
// ---------------------------------------------------------------------------

#[test]
fn test_cas() {
    let dir = tempdir().unwrap();
    let tree = open_counter_tree(dir.path());

    let c1 = Counters {
        views: 10,
        likes: 1,
    };
    tree.put(&1u64.to_be_bytes(), &c1).unwrap();

    let wrong = Counters {
        views: 999,
        likes: 0,
    };
    assert!(
        tree.cas(&1u64.to_be_bytes(), &wrong, &Counters::default())
            .is_err()
    );

    let c2 = Counters {
        views: 20,
        likes: 2,
    };
    tree.cas(&1u64.to_be_bytes(), &c1, &c2).unwrap();
    assert_eq!(tree.get(&1u64.to_be_bytes()).unwrap(), c2);
}

#[test]
fn test_update() {
    let dir = tempdir().unwrap();
    let tree = open_counter_tree(dir.path());

    tree.put(
        &1u64.to_be_bytes(),
        &Counters {
            views: 10,
            likes: 1,
        },
    )
    .unwrap();

    let new = tree
        .update(&1u64.to_be_bytes(), |c| Counters {
            views: c.views + 1,
            likes: c.likes,
        })
        .unwrap()
        .unwrap();
    assert_eq!(new.views, 11);
}

// ---------------------------------------------------------------------------
// Iteration
// ---------------------------------------------------------------------------

#[test]
fn test_iter() {
    let dir = tempdir().unwrap();
    let tree = open_counter_tree(dir.path());
    populate(&tree, 10);

    // reversed=true → descending
    let entries: Vec<_> = tree.iter().collect();
    assert_eq!(entries.len(), 10);

    for i in 0..10u64 {
        let expected = 9 - i;
        assert_eq!(entries[i as usize].0, expected.to_be_bytes());
        assert_eq!(entries[i as usize].1.views, expected * 10);
    }
}

#[test]
fn test_iter_rev() {
    let dir = tempdir().unwrap();
    let tree = open_counter_tree(dir.path());
    populate(&tree, 10);

    // .rev() on reversed=true → ascending
    let entries: Vec<_> = tree.iter().rev().collect();
    assert_eq!(entries.len(), 10);
    assert_eq!(entries[0].1.views, 0); // key=0, views=0
    assert_eq!(entries[9].1.views, 90); // key=9, views=90
}

#[test]
fn test_range() {
    let dir = tempdir().unwrap();
    let tree = open_counter_tree(dir.path());
    populate(&tree, 20);

    // reversed=true: range [5,10) → descending: 9,8,7,6,5
    let start = 5u64.to_be_bytes();
    let end = 10u64.to_be_bytes();
    let entries: Vec<_> = tree.range(&start, &end).collect();
    assert_eq!(entries.len(), 5);
    assert_eq!(entries[0].1.views, 90);
    assert_eq!(entries[4].1.views, 50);
}

#[test]
fn test_prefix_iter() {
    let dir = tempdir().unwrap();
    let tree: ZeroTree<[u8; 16], { size_of::<Counters>() }, Counters> =
        ZeroTree::<[u8; 16], { size_of::<Counters>() }, Counters>::open(dir.path(), Config::test())
            .unwrap();

    for i in 0u64..5 {
        let mut key = [0u8; 16];
        key[3] = 1;
        key[8..].copy_from_slice(&i.to_be_bytes());
        tree.put(&key, &Counters { views: i, likes: 0 }).unwrap();
    }
    for i in 0u64..3 {
        let mut key = [0u8; 16];
        key[3] = 2;
        key[8..].copy_from_slice(&i.to_be_bytes());
        tree.put(
            &key,
            &Counters {
                views: 100 + i,
                likes: 0,
            },
        )
        .unwrap();
    }

    assert_eq!(tree.prefix_iter(&[0, 0, 0, 1]).count(), 5);
    assert_eq!(tree.prefix_iter(&[0, 0, 0, 2]).count(), 3);
    assert_eq!(tree.prefix_iter(&[0, 0, 0, 3]).count(), 0);
}

// ---------------------------------------------------------------------------
// Reversed ordering
// ---------------------------------------------------------------------------

#[test]
fn test_reversed() {
    let dir = tempdir().unwrap();
    let mut config = Config::test();
    config.reversed = true;
    let tree: CounterTree = CounterTree::open(dir.path(), config).unwrap();
    populate(&tree, 10);

    let entries: Vec<_> = tree.iter().collect();
    assert_eq!(entries.len(), 10);
    // Forward on reversed = descending keys
    assert_eq!(entries[0].1.views, 90); // key=9
    assert_eq!(entries[9].1.views, 0); // key=0
}

// ---------------------------------------------------------------------------
// Different repr(C) type
// ---------------------------------------------------------------------------

#[test]
fn test_point_type() {
    let dir = tempdir().unwrap();
    type PointTree = ZeroTree<[u8; 8], { size_of::<Point>() }, Point>;
    let tree: PointTree = PointTree::open(dir.path(), Config::test()).unwrap();

    tree.put(
        &1u64.to_be_bytes(),
        &Point {
            x: 10,
            y: 20,
            z: 30,
        },
    )
    .unwrap();

    let p = tree.get(&1u64.to_be_bytes()).unwrap();
    assert_eq!(
        p,
        Point {
            x: 10,
            y: 20,
            z: 30
        }
    );
}

// ---------------------------------------------------------------------------
// Recovery
// ---------------------------------------------------------------------------

#[test]
fn test_recovery() {
    let dir = tempdir().unwrap();

    {
        let tree = open_counter_tree(dir.path());
        populate(&tree, 50);
        tree.close().unwrap();
    }

    let tree = open_counter_tree(dir.path());
    assert_eq!(tree.len(), 50);

    for i in 0..50u64 {
        let c = tree.get(&i.to_be_bytes()).unwrap();
        assert_eq!(c.views, i * 10);
        assert_eq!(c.likes, i);
    }
}

#[test]
fn test_recovery_with_deletes() {
    let dir = tempdir().unwrap();

    {
        let tree = open_counter_tree(dir.path());
        populate(&tree, 20);
        for i in (0..20u64).step_by(2) {
            tree.delete(&i.to_be_bytes()).unwrap();
        }
        tree.close().unwrap();
    }

    let tree = open_counter_tree(dir.path());
    assert_eq!(tree.len(), 10);

    for i in 0..20u64 {
        if i % 2 == 1 {
            assert!(tree.contains(&i.to_be_bytes()));
        } else {
            assert!(!tree.contains(&i.to_be_bytes()));
        }
    }
}

// ---------------------------------------------------------------------------
// Migrate
// ---------------------------------------------------------------------------

#[test]
fn test_migrate() {
    let dir = tempdir().unwrap();
    let tree = open_counter_tree(dir.path());
    populate(&tree, 10);

    let mutations = tree
        .migrate(|_key, c| {
            if c.likes < 5 {
                MigrateAction::Update(Counters {
                    views: c.views + 1000,
                    likes: c.likes,
                })
            } else {
                MigrateAction::Keep
            }
        })
        .unwrap();

    assert_eq!(mutations, 5);
    assert_eq!(tree.get(&0u64.to_be_bytes()).unwrap().views, 1000);
    assert_eq!(tree.get(&5u64.to_be_bytes()).unwrap().views, 50); // unchanged
}

// ---------------------------------------------------------------------------
// Compact
// ---------------------------------------------------------------------------

#[test]
fn test_compact() {
    let dir = tempdir().unwrap();
    let tree = open_counter_tree(dir.path());

    populate(&tree, 100);
    // Overwrite all
    for i in 0..100u64 {
        tree.put(
            &i.to_be_bytes(),
            &Counters {
                views: i * 100,
                likes: 0,
            },
        )
        .unwrap();
    }

    tree.flush_buffers().unwrap();
    tree.compact().unwrap();

    for i in 0..100u64 {
        assert_eq!(tree.get(&i.to_be_bytes()).unwrap().views, i * 100);
    }
}