use shardize::shardize;

use std::collections::HashMap;

// This trait, its impl for HashMap and the shard_key are the minimum
// required boilerplate to make a sharded container.

#[shardize(ShardedHashMap)]
trait MyTrait: Default {
    fn get(&self, key: usize) -> Option<&String>;
    fn insert(&mut self, key: usize, value: String);
}

impl MyTrait for HashMap<usize, String> {
    fn get(&self, key: usize) -> Option<&String> {
        (self as &HashMap<usize, String>).get(&key)
    }

    fn insert(&mut self, key: usize, value: String) {
        (self as &mut HashMap<usize, String>).insert(key, value);
    }
}

fn my_shard_key(key: usize) -> usize {
    key
}

// Optional: demonstrates that you can do stuff across shards if you want to
impl<const NUM_SHARDS: usize>
    ShardedHashMap<HashMap<usize, String>, NUM_SHARDS>
{
    fn values(&self) -> Vec<String> {
        self.shards
            .iter()
            .flat_map(|shard| shard.values())
            .cloned()
            .collect()
    }
}

#[test]
fn can_get_and_set_values() {
    let mut m = ShardedHashMap::<HashMap<usize, String>, 3>::new(&my_shard_key);

    m.insert(5, String::from("foo"));
    assert_eq!(m.get(5).unwrap(), "foo");
}

#[test]
fn values_returns_all_values() {
    let mut m = ShardedHashMap::<HashMap<usize, String>, 3>::new(&my_shard_key);

    m.insert(5, String::from("foo"));
    m.insert(6, String::from("bar"));

    let mut actual = m.values();
    let mut expected = vec![String::from("foo"), String::from("bar")];
    actual.sort();
    expected.sort();

    assert_eq!(actual, expected);
}