use armdb::{Config, ConstTree, DbError};
use tempfile::tempdir;

#[test]
fn test_prefix_bits_same_prefix_same_shard() {
    let dir = tempdir().unwrap();
    let mut config = Config::test();
    config.shard_prefix_bits = 32;
    config.shard_count = 8;

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

    // Keys with identical first 4 bytes should route to the same shard
    let prefix = [0x01, 0x02, 0x03, 0x04];
    let mut keys = Vec::new();
    for i in 0..10u32 {
        let mut key = [0u8; 8];
        key[..4].copy_from_slice(&prefix);
        key[4..].copy_from_slice(&i.to_be_bytes());
        keys.push(key);
    }

    let first_shard = tree.shard_for(&keys[0]);
    for key in &keys[1..] {
        assert_eq!(
            tree.shard_for(key),
            first_shard,
            "keys with same prefix should route to same shard"
        );
    }
}

#[test]
fn test_prefix_bits_different_prefix_may_differ() {
    let dir = tempdir().unwrap();
    let mut config = Config::test();
    config.shard_prefix_bits = 32;
    config.shard_count = 8;

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

    // Keys with different first 4 bytes — at least some should route to different shards
    let mut shards = std::collections::HashSet::new();
    for p in 0..20u32 {
        let mut key = [0u8; 8];
        key[..4].copy_from_slice(&p.to_be_bytes());
        shards.insert(tree.shard_for(&key));
    }

    assert!(
        shards.len() > 1,
        "20 different prefixes across 8 shards should hit more than 1 shard, got {}",
        shards.len()
    );
}

#[test]
fn test_atomic_same_shard_succeeds() {
    let dir = tempdir().unwrap();
    let mut config = Config::test();
    config.shard_prefix_bits = 32;
    config.shard_count = 4;

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

    // Two keys with same prefix
    let key_a = [0x01, 0x02, 0x03, 0x04, 0x00, 0x00, 0x00, 0x01];
    let key_b = [0x01, 0x02, 0x03, 0x04, 0x00, 0x00, 0x00, 0x02];
    assert_eq!(tree.shard_for(&key_a), tree.shard_for(&key_b));

    let val_a = 10u64.to_be_bytes();
    let val_b = 20u64.to_be_bytes();

    tree.put(&key_a, &1u64.to_be_bytes()).unwrap();
    tree.put(&key_b, &2u64.to_be_bytes()).unwrap();

    // Atomic operation on both keys in one shard
    tree.atomic(&key_a, |shard| {
        shard.put(&key_a, &val_a)?;
        shard.put(&key_b, &val_b)?;
        Ok(())
    })
    .unwrap();

    assert_eq!(tree.get(&key_a).unwrap(), val_a);
    assert_eq!(tree.get(&key_b).unwrap(), val_b);
}

#[test]
fn test_atomic_different_shard_fails() {
    let dir = tempdir().unwrap();
    let mut config = Config::test();
    config.shard_prefix_bits = 32;
    config.shard_count = 8;

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

    // Find two keys that route to different shards
    let key_a = [0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01];
    let mut key_b = None;
    for p in 1..100u32 {
        let mut candidate = [0u8; 8];
        candidate[..4].copy_from_slice(&p.to_be_bytes());
        candidate[7] = 1;
        if tree.shard_for(&candidate) != tree.shard_for(&key_a) {
            key_b = Some(candidate);
            break;
        }
    }
    let key_b = key_b.expect("should find a key routing to a different shard");

    tree.put(&key_a, &1u64.to_be_bytes()).unwrap();
    tree.put(&key_b, &2u64.to_be_bytes()).unwrap();

    // Atomic: key_b routes to a different shard → ShardMismatch
    let result = tree.atomic(&key_a, |shard| {
        shard.put(&key_b, &99u64.to_be_bytes())?;
        Ok(())
    });

    assert!(
        matches!(result, Err(DbError::ShardMismatch)),
        "expected ShardMismatch, got {:?}",
        result
    );
}

#[test]
fn test_prefix_bits_iteration() {
    let dir = tempdir().unwrap();
    let mut config = Config::test();
    config.shard_prefix_bits = 32;

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

    // 3 groups of 10 entries with different prefixes
    for group in 0..3u32 {
        for id in 0..10u32 {
            let mut key = [0u8; 8];
            key[..4].copy_from_slice(&group.to_be_bytes());
            key[4..].copy_from_slice(&id.to_be_bytes());
            let val = ((group as u64) * 100 + id as u64).to_be_bytes();
            tree.put(&key, &val).unwrap();
        }
    }

    assert_eq!(tree.len(), 30);

    // prefix_iter for each group
    for group in 0..3u32 {
        let prefix = group.to_be_bytes();
        let entries: Vec<_> = tree.prefix_iter(&prefix).collect();
        assert_eq!(
            entries.len(),
            10,
            "group {} should have 10 entries, got {}",
            group,
            entries.len()
        );
    }

    // Full iteration returns all 30
    let all: Vec<_> = tree.iter().collect();
    assert_eq!(all.len(), 30);
}

#[test]
fn test_prefix_bits_zero_default() {
    let dir = tempdir().unwrap();
    let config = Config::test(); // shard_prefix_bits = 0 by default

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

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

    assert_eq!(tree.len(), 100);

    for i in 0..100u64 {
        let key = i.to_be_bytes();
        let expected = (i * 10).to_be_bytes();
        assert_eq!(tree.get(&key).unwrap(), expected);
    }
}
