//! A sharded concurrent map for reducing lock contention.
//!
//! This module provides [`ShardedMap`], a map that distributes entries across
//! `N` shards, each protected by its own mutex. This reduces contention compared
//! to a single global lock, especially under high concurrency.
//!
//! Shard selection uses SipHash-2-4 with a random key to prevent `DoS` attacks
//! where an adversary crafts keys that all hash to the same shard.

use alloc::vec::Vec;
use async_lock::Mutex;
use core::hash::{Hash, Hasher};
use sedimentree_core::collections::Map;
use siphasher::sip::SipHasher24;

/// A sharded concurrent map with `N` independent shards.
///
/// Each shard contains a separate [`Map`] protected by its own [`Mutex`],
/// allowing concurrent access to different shards. Shard selection uses
/// SipHash-2-4 with a secret key for `DoS` resistance.
///
/// # Type Parameters
///
/// - `K`: The key type, must implement [`Hash`]
/// - `V`: The value type
/// - `N`: The number of shards (default: 256)
///
/// # Example
///
/// ```ignore
/// use subduction_core::sharded_map::ShardedMap;
///
/// // With random key (requires `getrandom` feature)
/// let map: ShardedMap<SedimentreeId, Sedimentree> = ShardedMap::new();
///
/// // With explicit key (deterministic, for testing)
/// let map: ShardedMap<SedimentreeId, Sedimentree> = ShardedMap::with_key(0, 0);
/// ```
#[derive(Debug)]
pub struct ShardedMap<K: Hash + Ord, V, const N: usize = 256> {
    shards: [Mutex<Map<K, V>>; N],

    /// Seed 0 for keyed hash.
    key0: u64,

    /// Seed 1 for keyed hash.
    key1: u64,
}

impl<K: Hash + Ord, V, const N: usize> ShardedMap<K, V, N> {
    /// Creates a new empty [`ShardedMap`] with a randomly generated key.
    ///
    /// # Panics
    ///
    /// Panics if the system random number generator fails.
    #[must_use]
    pub fn new() -> Self {
        let mut key_bytes = [0u8; 16];

        #[allow(clippy::expect_used)]
        getrandom::getrandom(&mut key_bytes).expect("getrandom failed");

        #[allow(clippy::expect_used)]
        let key0 = u64::from_le_bytes(
            key_bytes[0..8]
                .try_into()
                .expect("generated too few key bytes"),
        );

        #[allow(clippy::expect_used)]
        let key1 = u64::from_le_bytes(
            key_bytes[8..16]
                .try_into()
                .expect("generated too few key bytes"),
        );

        Self::with_key(key0, key1)
    }

    /// Creates a new empty [`ShardedMap`] with the given `SipHash` keys.
    ///
    /// Use this constructor when you need deterministic behavior (e.g., testing)
    /// or when you want to persist and restore the key.
    #[must_use]
    pub fn with_key(key0: u64, key1: u64) -> Self {
        Self {
            shards: core::array::from_fn(|_| Mutex::new(Map::new())),
            key0,
            key1,
        }
    }

    /// Returns the `SipHash` keys used for shard selection.
    ///
    /// Useful for persisting the key to restore deterministic behavior.
    #[must_use]
    pub const fn sip_keys(&self) -> (u64, u64) {
        (self.key0, self.key1)
    }

    /// Returns the shard index for the given key.
    ///
    /// # Note
    ///
    /// Uses [Lemire's "fast range" method][post] to map hash to $[0, N)$ without modulo bias.
    ///
    /// [post]: https://lemire.me/blog/2016/06/27/a-fast-alternative-to-the-modulo-reduction
    #[inline]
    fn shard_index(&self, key: &K) -> usize {
        let mut hasher = SipHasher24::new_with_keys(self.key0, self.key1);
        key.hash(&mut hasher);
        let hash = hasher.finish();

        #[allow(clippy::expect_used)]
        usize::try_from((u128::from(hash) * (N as u128)) >> 64).expect("N must stay in usize")
    }

    /// Returns a reference to the shard mutex at the given index, if valid.
    ///
    /// Returns `None` if `index >= N`.
    #[inline]
    pub const fn shard_at(&self, index: usize) -> Option<&Mutex<Map<K, V>>> {
        if index < N {
            #[allow(clippy::indexing_slicing)] // checked above + no `get` on array to avoid nighly
            Some(&self.shards[index])
        } else {
            None
        }
    }

    /// Returns a reference to the shard mutex for the given key.
    ///
    /// This is useful for complex operations that need to hold the lock
    /// across multiple operations or async boundaries.
    ///
    /// # Example
    ///
    /// ```ignore
    /// let shard = map.get_shard_containing(&key);
    /// let mut guard = shard.lock().await;
    /// // Complex operations with the guard...
    /// ```
    #[inline]
    pub fn get_shard_containing(&self, key: &K) -> &Mutex<Map<K, V>> {
        #[allow(clippy::indexing_slicing)] // shard_index always returns valid index
        &self.shards[self.shard_index(key)]
    }

    /// Returns the number of shards.
    #[inline]
    #[must_use]
    pub const fn shard_count(&self) -> usize {
        N
    }

    /// Gets a cloned value for the given key, if it exists.
    pub async fn get_cloned(&self, key: &K) -> Option<V>
    where
        V: Clone,
    {
        self.get_shard_containing(key)
            .lock()
            .await
            .get(key)
            .cloned()
    }

    /// Returns `true` if the map contains the given key.
    pub async fn contains_key(&self, key: &K) -> bool {
        self.get_shard_containing(key)
            .lock()
            .await
            .contains_key(key)
    }

    /// Inserts a key-value pair into the map.
    ///
    /// Returns the previous value if the key was already present.
    pub async fn insert(&self, key: K, value: V) -> Option<V> {
        self.get_shard_containing(&key)
            .lock()
            .await
            .insert(key, value)
    }

    /// Removes a key from the map, returning the value if it was present.
    pub async fn remove(&self, key: &K) -> Option<V> {
        self.get_shard_containing(key).lock().await.remove(key)
    }

    /// Gets or inserts a default value for the given key.
    ///
    /// Returns a clone of the value (either existing or newly inserted).
    pub async fn entry_or_default(&self, key: K) -> V
    where
        V: Default + Clone,
    {
        self.get_shard_containing(&key)
            .lock()
            .await
            .entry(key)
            .or_default()
            .clone()
    }

    /// Gets or inserts a default value, then applies a function to it.
    ///
    /// This is useful for mutating the value in place without cloning.
    pub async fn with_entry_or_default<F: FnOnce(&mut V) -> R, R>(&self, key: K, f: F) -> R
    where
        V: Default,
    {
        let mut guard = self.get_shard_containing(&key).lock().await;
        let value = guard.entry(key).or_default();
        f(value)
    }

    /// Applies a function to a value if it exists.
    ///
    /// Returns `None` if the key doesn't exist, otherwise returns the result of `f`.
    pub async fn with_entry<F: FnOnce(&mut V) -> R, R>(&self, key: &K, f: F) -> Option<R> {
        let mut guard = self.get_shard_containing(key).lock().await;
        guard.get_mut(key).map(f)
    }

    /// Collects all keys from all shards.
    ///
    /// Note: This acquires each shard lock sequentially, releasing between shards.
    /// For very large maps, consider using [`shard_at`] for incremental processing.
    pub async fn into_keys(&self) -> Vec<K>
    where
        K: Clone,
    {
        let mut keys = Vec::new();
        for shard in &self.shards {
            keys.extend(shard.lock().await.keys().cloned());
        }
        keys
    }

    /// Collects all values from all shards.
    ///
    /// Note: This acquires each shard lock sequentially, releasing between shards.
    pub async fn into_values(&self) -> Vec<V>
    where
        V: Clone,
    {
        let mut values = Vec::new();
        for shard in &self.shards {
            values.extend(shard.lock().await.values().cloned());
        }
        values
    }

    /// Returns an iterator over shard indices.
    ///
    /// Useful for shard-by-shard iteration where you want explicit control
    /// over lock acquisition and release:
    ///
    /// ```ignore
    /// for idx in map.shard_indices() {
    ///     if let Some(shard) = map.shard_at(idx) {
    ///         let guard = shard.lock().await;
    ///         // Process shard, lock is released at end of scope
    ///     }
    /// }
    /// ```
    pub fn shard_indices(&self) -> impl Iterator<Item = usize> {
        0..N
    }

    /// Merges a value into the map using a merge function.
    ///
    /// If the key exists, calls `merge(existing, value)`.
    /// If the key doesn't exist, inserts the value directly.
    pub async fn merge_with<F: FnOnce(&mut V, V)>(&self, key: K, value: V, merge: F) {
        let mut guard = self.get_shard_containing(&key).lock().await;
        match guard.get_mut(&key) {
            Some(existing) => merge(existing, value),
            None => {
                guard.insert(key, value);
            }
        }
    }

    /// Returns the total number of entries across all shards.
    ///
    /// Note: This acquires each shard lock sequentially.
    pub async fn len(&self) -> usize {
        let mut total = 0;
        for shard in &self.shards {
            total += shard.lock().await.len();
        }
        total
    }

    /// Returns `true` if the map is empty.
    ///
    /// Note: This may acquire multiple shard locks.
    pub async fn is_empty(&self) -> bool {
        for shard in &self.shards {
            if !shard.lock().await.is_empty() {
                return false;
            }
        }
        true
    }

    /// Clears all entries from all shards.
    pub async fn clear(&self) {
        for shard in &self.shards {
            shard.lock().await.clear();
        }
    }
}

impl<K: Hash + Ord, V, const N: usize> Default for ShardedMap<K, V, N> {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use alloc::{collections::BTreeSet, vec};
    use sedimentree_core::id::SedimentreeId;
    use testresult::TestResult;

    #[tokio::test]
    async fn test_basic_operations() {
        let map: ShardedMap<SedimentreeId, u32, 16> = ShardedMap::with_key(0, 0);

        let id1 = SedimentreeId::new([1; 32]);
        let id2 = SedimentreeId::new([2; 32]);

        // Insert
        assert!(map.insert(id1, 100).await.is_none());
        assert!(map.insert(id2, 200).await.is_none());

        // Get
        assert_eq!(map.get_cloned(&id1).await, Some(100));
        assert_eq!(map.get_cloned(&id2).await, Some(200));

        // Contains
        assert!(map.contains_key(&id1).await);
        assert!(!map.contains_key(&SedimentreeId::new([3; 32])).await);

        // Remove
        assert_eq!(map.remove(&id1).await, Some(100));
        assert!(map.get_cloned(&id1).await.is_none());

        // Length
        assert_eq!(map.len().await, 1);
    }

    #[cfg(all(test, feature = "std", feature = "bolero"))]
    mod proptests {
        use super::*;

        #[test]
        fn prop_shard_distribution_uses_multiple_shards() {
            bolero::check!()
                .with_arbitrary::<(u64, u64, [u8; 32], [u8; 32], [u8; 32], [u8; 32])>()
                .for_each(|(key0, key1, id0, id1, id2, id3)| {
                    let map: ShardedMap<SedimentreeId, (), 16> = ShardedMap::with_key(*key0, *key1);

                    // With 4 random IDs into 16 shards, we should usually see multiple shards used
                    let ids = [
                        SedimentreeId::new(*id0),
                        SedimentreeId::new(*id1),
                        SedimentreeId::new(*id2),
                        SedimentreeId::new(*id3),
                    ];

                    let mut shards_used = BTreeSet::new();
                    for id in &ids {
                        shards_used.insert(map.shard_index(id));
                    }

                    // At minimum, not all IDs should hash to the same shard (except in rare cases)
                    // This is a probabilistic test - with 4 items and 16 shards, probability of
                    // all same shard is (1/16)^3 ≈ 0.02%, so this should almost always pass
                    // We're just checking basic sanity that hashing works
                    assert!(!shards_used.is_empty(), "at least one shard should be used");
                });
        }

        #[test]
        fn prop_different_keys_give_different_shard_indices() {
            bolero::check!()
                .with_arbitrary::<(u64, u64, u64, u64, [u8; 32])>()
                .for_each(|(key0a, key1a, key0b, key1b, id_bytes)| {
                    // Skip if keys are the same
                    if *key0a == *key0b && *key1a == *key1b {
                        return;
                    }

                    let map_a: ShardedMap<SedimentreeId, (), 256> =
                        ShardedMap::with_key(*key0a, *key1a);
                    let map_b: ShardedMap<SedimentreeId, (), 256> =
                        ShardedMap::with_key(*key0b, *key1b);

                    let id = SedimentreeId::new(*id_bytes);

                    let idx_a = map_a.shard_index(&id);
                    let idx_b = map_b.shard_index(&id);

                    // Both should be valid shard indices
                    assert!(idx_a < 256);
                    assert!(idx_b < 256);

                    // Note: we don't assert they're different because with 256 shards
                    // there's a 1/256 chance of collision even with different keys
                });
        }
    }

    #[tokio::test]
    async fn test_with_entry_or_default() -> TestResult {
        let map: ShardedMap<SedimentreeId, Vec<u32>, 16> = ShardedMap::with_key(0, 0);
        let id = SedimentreeId::new([1; 32]);

        map.with_entry_or_default(id, |v| v.push(1)).await;
        map.with_entry_or_default(id, |v| v.push(2)).await;

        let value = map.get_cloned(&id).await.ok_or("should exist")?;
        assert_eq!(value, vec![1, 2]);

        Ok(())
    }

    #[tokio::test]
    async fn test_merge_with() {
        let map: ShardedMap<SedimentreeId, u32, 16> = ShardedMap::with_key(0, 0);
        let id = SedimentreeId::new([1; 32]);

        map.merge_with(id, 10, |existing, new| *existing += new)
            .await;
        assert_eq!(map.get_cloned(&id).await, Some(10));

        map.merge_with(id, 5, |existing, new| *existing += new)
            .await;
        assert_eq!(map.get_cloned(&id).await, Some(15));
    }
}
