use dashmap::DashMap;
use std::{
    collections::BTreeSet,
    hash::Hash,
    mem::ManuallyDrop,
    ops::{Deref, DerefMut},
    sync::Arc,
};
use tokio::sync::{OwnedRwLockReadGuard, OwnedRwLockWriteGuard, RwLock};

pub trait Empty {
    fn is_empty(&self) -> bool;
}

impl Empty for () {
    fn is_empty(&self) -> bool {
        true
    }
}

impl<V> Empty for BTreeSet<V> {
    fn is_empty(&self) -> bool {
        BTreeSet::is_empty(self)
    }
}

type KeyLockInner<K, V> = Arc<DashMap<K, Arc<RwLock<V>>>>;

#[derive(Clone)]
pub struct KeyLocks<K, V>(KeyLockInner<K, V>);

struct KeyRef<K, V> {
    map: KeyLockInner<K, V>,
    key: K,
}
impl<K: Eq + Hash, V: Empty> KeyRef<K, V> {
    fn remove_if_needed<G: Deref<Target = V>>(&self, guard: &mut ManuallyDrop<G>) {
        if guard.is_empty() {
            self.map.remove_if(&self.key, |_, v| {
                unsafe { ManuallyDrop::drop(guard) };
                Arc::strong_count(v) == 1
            });
        } else {
            unsafe { ManuallyDrop::drop(guard) };
        }
    }
}

pub struct OwnedReadGuard<K: Eq + Hash, V: Empty> {
    key: KeyRef<K, V>,
    guard: ManuallyDrop<OwnedRwLockReadGuard<V>>,
}
impl<K: Eq + Hash, V: Empty> Deref for OwnedReadGuard<K, V> {
    type Target = V;
    fn deref(&self) -> &Self::Target {
        self.guard.deref()
    }
}
impl<K: Eq + Hash, V: Empty> Drop for OwnedReadGuard<K, V> {
    fn drop(&mut self) {
        self.key.remove_if_needed(&mut self.guard);
    }
}

pub struct OwnedWriteGuard<K: Eq + Hash, V: Empty> {
    key: KeyRef<K, V>,
    guard: ManuallyDrop<OwnedRwLockWriteGuard<V>>,
}
impl<K: Eq + Hash, V: Empty> Deref for OwnedWriteGuard<K, V> {
    type Target = V;
    fn deref(&self) -> &Self::Target {
        self.guard.deref()
    }
}

impl<K: Eq + Hash, V: Empty> DerefMut for OwnedWriteGuard<K, V> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.guard.deref_mut()
    }
}

impl<K: Eq + Hash, V: Empty> Drop for OwnedWriteGuard<K, V> {
    fn drop(&mut self) {
        self.key.remove_if_needed(&mut self.guard);
    }
}

impl<K: Eq + Hash + Clone, V: Empty + Default> KeyLocks<K, V> {
    pub fn new() -> Self {
        KeyLocks(Arc::default())
    }

    fn key_ref(&self, key: K) -> KeyRef<K, V> {
        KeyRef {
            map: self.0.clone(),
            key,
        }
    }

    pub async fn read(&self, key: K) -> OwnedReadGuard<K, V> {
        OwnedReadGuard {
            key: self.key_ref(key.clone()),
            guard: ManuallyDrop::new(
                self.0
                    .entry(key)
                    .or_insert_with(Arc::default)
                    .clone()
                    .read_owned()
                    .await,
            ),
        }
    }

    pub async fn write(&self, key: K) -> OwnedWriteGuard<K, V> {
        OwnedWriteGuard {
            key: self.key_ref(key.clone()),
            guard: ManuallyDrop::new(
                self.0
                    .entry(key)
                    .or_insert_with(Arc::default)
                    .clone()
                    .write_owned()
                    .await,
            ),
        }
    }

    pub fn len(&self) -> usize {
        self.0.len()
    }
}

#[cfg(test)]
mod test {
    use super::KeyLocks;
    use std::collections::BTreeSet;

    #[tokio::test]
    async fn drop_only_if_empty() {
        let locks = KeyLocks::<u32, BTreeSet<String>>::new();

        let mut lock = locks.write(1).await;
        lock.insert("Hello".to_owned());
        lock.insert("World".to_owned());
        drop(lock);

        // Value is not empty and thus is not dropped
        assert_eq!(locks.len(), 1);

        let mut lock = locks.write(1).await;
        assert_eq!(lock.len(), 2);
        lock.clear();
        drop(lock);

        // Should be dropped now
        assert_eq!(locks.len(), 0);
    }
}