#[cfg(feature = "var-collections")]
use armdb::{ByteView, VarTree};
use armdb::{Config, ConstTree, MigrateAction};
use tempfile::tempdir;

#[test]
fn test_migrate_large_const_tree() {
    let dir = tempdir().unwrap();
    let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), Config::test()).unwrap();

    // Write 10_000 entries
    for i in 0..10_000u64 {
        let key = i.to_be_bytes();
        let val = i.to_be_bytes();
        tree.put(&key, &val).unwrap();
    }

    assert_eq!(tree.len(), 10_000);

    // Migrate: update first 5_000, keep the rest
    let mutations = tree
        .migrate(|key, _val| {
            let i = u64::from_be_bytes(*key);
            if i < 5_000 {
                MigrateAction::Update((i + 100_000).to_be_bytes())
            } else {
                MigrateAction::Keep
            }
        })
        .unwrap();

    assert_eq!(mutations, 5_000);
    assert_eq!(tree.len(), 10_000);

    // Verify updated entries
    for i in 0..5_000u64 {
        let key = i.to_be_bytes();
        let expected = (i + 100_000).to_be_bytes();
        assert_eq!(
            tree.get(&key).unwrap(),
            expected,
            "updated value mismatch at key {}",
            i
        );
    }

    // Verify kept entries
    for i in 5_000..10_000u64 {
        let key = i.to_be_bytes();
        let expected = i.to_be_bytes();
        assert_eq!(
            tree.get(&key).unwrap(),
            expected,
            "kept value mismatch at key {}",
            i
        );
    }
}

#[test]
fn test_migrate_large_delete() {
    let dir = tempdir().unwrap();
    let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), Config::test()).unwrap();

    for i in 0..10_000u64 {
        let key = i.to_be_bytes();
        let val = i.to_be_bytes();
        tree.put(&key, &val).unwrap();
    }

    // Delete every other entry
    let mutations = tree
        .migrate(|key, _val| {
            let i = u64::from_be_bytes(*key);
            if i % 2 == 0 {
                MigrateAction::Delete
            } else {
                MigrateAction::Keep
            }
        })
        .unwrap();

    assert_eq!(mutations, 5_000);
    assert_eq!(tree.len(), 5_000);

    // Verify even keys are gone, odd keys remain
    for i in 0..10_000u64 {
        let key = i.to_be_bytes();
        if i % 2 == 0 {
            assert!(tree.get(&key).is_none(), "key {} should be deleted", i);
        } else {
            assert_eq!(tree.get(&key).unwrap(), i.to_be_bytes());
        }
    }
}

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

    {
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), Config::test()).unwrap();

        for i in 0..5_000u64 {
            let key = i.to_be_bytes();
            let val = i.to_be_bytes();
            tree.put(&key, &val).unwrap();
        }

        // Migrate: multiply values by 10
        let mutations = tree
            .migrate(|_key, val| {
                let v = u64::from_be_bytes(*val);
                MigrateAction::Update((v * 10).to_be_bytes())
            })
            .unwrap();
        assert_eq!(mutations, 5_000);

        tree.close().unwrap();
    }

    // Reopen and verify migrated values survived recovery
    {
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), Config::test()).unwrap();

        assert_eq!(tree.len(), 5_000);

        for i in 0..5_000u64 {
            let key = i.to_be_bytes();
            let expected = (i * 10).to_be_bytes();
            assert_eq!(
                tree.get(&key).unwrap(),
                expected,
                "recovered value mismatch at key {}",
                i
            );
        }
    }
}

#[cfg(feature = "var-collections")]
#[test]
fn test_migrate_large_var_tree() {
    let dir = tempdir().unwrap();
    let tree = VarTree::<[u8; 8]>::open(dir.path(), Config::test()).unwrap();

    for i in 0..10_000u64 {
        let key = i.to_be_bytes();
        let val = format!("val_{i}");
        tree.put(&key, val.as_bytes()).unwrap();
    }

    assert_eq!(tree.len(), 10_000);

    // Update first 5_000 with longer values
    let mutations = tree
        .migrate(|key, _val| {
            let i = u64::from_be_bytes(*key);
            if i < 5_000 {
                let new_val = format!("updated_{i}");
                MigrateAction::Update(ByteView::new(new_val.as_bytes()))
            } else {
                MigrateAction::Keep
            }
        })
        .unwrap();

    assert_eq!(mutations, 5_000);
    assert_eq!(tree.len(), 10_000);

    // Verify
    for i in 0..5_000u64 {
        let key = i.to_be_bytes();
        let expected = format!("updated_{i}");
        assert_eq!(tree.get(&key).unwrap().as_ref(), expected.as_bytes());
    }
    for i in 5_000..10_000u64 {
        let key = i.to_be_bytes();
        let expected = format!("val_{i}");
        assert_eq!(tree.get(&key).unwrap().as_ref(), expected.as_bytes());
    }
}

#[test]
fn test_migrate_all_delete() {
    let dir = tempdir().unwrap();
    let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), Config::test()).unwrap();

    for i in 0..1_000u64 {
        let key = i.to_be_bytes();
        let val = i.to_be_bytes();
        tree.put(&key, &val).unwrap();
    }

    let mutations = tree.migrate(|_, _| MigrateAction::Delete).unwrap();
    assert_eq!(mutations, 1_000);
    assert_eq!(tree.len(), 0);

    for i in 0..1_000u64 {
        let key = i.to_be_bytes();
        assert!(tree.get(&key).is_none());
    }
}
