use core::{
    borrow::Borrow,
    hash::{BuildHasher, Hash},
};

use alloc::vec::Vec;

pub type DefaultHashBuilder = core::hash::BuildHasherDefault<ahash::AHasher>;
pub type DefaultHashTable<K, V, const N: usize = 32> = HashTable<K, V, DefaultHashBuilder, N>;

pub struct HashTable<K, V, H = DefaultHashBuilder, const N: usize = 32>
where
    [Vec<(K, V)>; N]: Sized,
{
    buckets: [Vec<(K, V)>; N],
    hasher: H,
}

impl<K, V, const N: usize> HashTable<K, V, DefaultHashBuilder, N>
where
    [Vec<(K, V)>; N]: Sized,
{
    pub const fn new() -> Self {
        Self {
            buckets: [const { Vec::new() }; N],
            hasher: DefaultHashBuilder::new(),
        }
    }
}

impl<K, V, H, const N: usize> HashTable<K, V, H, N> {
    pub fn clear(&mut self) {
        for bucket in self.buckets.iter_mut() {
            bucket.clear();
        }
    }
}

impl<K: Hash + Eq, V, const N: usize> HashTable<K, V, DefaultHashBuilder, N> {
    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
        let h = self.hasher.hash_one(&key);
        let bucket = &mut self.buckets[h as usize % self.buckets.len()];

        if let Some((_, slot)) = bucket.iter_mut().find(|(k, _)| k == &key) {
            Some(core::mem::replace(slot, value))
        } else {
            bucket.push((key, value));
            None
        }
    }

    pub fn remove<Q>(&mut self, key: &Q) -> Option<V>
    where
        Q: Hash + Eq + ?Sized,
        K: Borrow<Q>,
    {
        let h = self.hasher.hash_one(key);
        let bucket = &mut self.buckets[h as usize % self.buckets.len()];

        for i in 0..bucket.len() {
            if bucket[i].0.borrow() == key {
                let (_, value) = bucket.remove(i);
                return Some(value);
            }
        }

        None
    }

    pub fn get<Q>(&self, key: &Q) -> Option<&V>
    where
        Q: Hash + Eq + ?Sized,
        K: Borrow<Q>,
    {
        self.get_key_value(key).map(|p| p.1)
    }

    pub fn get_key_value<Q>(&self, key: &Q) -> Option<(&K, &V)>
    where
        Q: Hash + Eq + ?Sized,
        K: Borrow<Q>,
    {
        let h = self.hasher.hash_one(key);
        let bucket = &self.buckets[h as usize % self.buckets.len()];

        bucket
            .iter()
            .find_map(|(k, v)| (k.borrow() == key).then_some((k, v)))
    }

    pub fn contains_key<Q>(&self, key: &Q) -> bool
    where
        Q: Hash + Eq + ?Sized,
        K: Borrow<Q>,
    {
        let h = self.hasher.hash_one(key);
        let bucket = &self.buckets[h as usize % self.buckets.len()];

        bucket.iter().any(|(k, _)| (k.borrow() == key))
    }
}
