#![forbid(unsafe_code)]

pub trait Key: Copy + PartialEq {
    type Value;
    type Owner: Owner<Key = Self>;

    fn owner(self) -> Self::Owner;

    fn map<R>(self, f: impl FnOnce(&Self::Value) -> R) -> R;
    fn map_mut<R>(self, f: impl FnOnce(&mut Self::Value) -> R) -> R;
}

pub trait Owner: Clone + PartialEq {
    type Key: Key;
    fn key(&self) -> Self::Key;

    fn shared_count(&self) -> usize;
}

pub trait Store<T> {
    type Key: Key<Value = T>;

    fn insert(&self, value: T) -> <Self::Key as Key>::Owner;
}

pub mod local {
    use std::marker::PhantomData;

    /// Key is !Send + !Sync.
    /// ```compile_fail
    /// # fn test(v: hooks_gen::local::Key<()>) -> impl Send { v };
    /// ```
    /// ```compile_fail
    /// # fn test(v: hooks_gen::local::Key<()>) -> impl Sync { v };
    /// ```
    pub struct Key<T> {
        index: usize,
        _phantom: PhantomData<*const T>,
    }

    impl<T> Eq for Key<T> {}
    impl<T> PartialEq for Key<T> {
        fn eq(&self, other: &Self) -> bool {
            self.index == other.index
        }
    }

    impl<T> Key<T> {
        const fn new(index: usize) -> Self {
            Self {
                index,
                _phantom: PhantomData,
            }
        }
    }

    impl<T> Copy for Key<T> {}
    impl<T> Clone for Key<T> {
        fn clone(&self) -> Self {
            *self
        }
    }

    impl<T: 'static> super::Key for Key<T> {
        type Value = T;

        type Owner = Owner<T>;

        fn owner(self) -> Self::Owner {
            Store::register_owner_at(self)
        }

        fn map<R>(self, f: impl FnOnce(&Self::Value) -> R) -> R {
            Store::map_at(self, f)
        }

        /// Note that this locks the whole store.
        fn map_mut<R>(self, f: impl FnOnce(&mut Self::Value) -> R) -> R {
            Store::map_mut_at(self, f)
        }
    }

    pub struct Owner<T: 'static> {
        key: Key<T>,
    }

    impl<T: 'static> PartialEq for Owner<T> {
        fn eq(&self, other: &Self) -> bool {
            self.key == other.key
        }
    }

    pub struct Store;

    trait StoreImpl {
        fn insert<T: 'static>(value: T) -> Owner<T>;
        fn unregister_owner<T>(owner: &mut Owner<T>);
        fn register_owner_at<T>(key: Key<T>) -> Owner<T>;

        fn owners_count_at<T: 'static>(key: Key<T>) -> usize;

        fn map_at<T: 'static, R>(key: Key<T>, f: impl FnOnce(&T) -> R) -> R;
        fn map_mut_at<T: 'static, R>(key: Key<T>, f: impl FnOnce(&mut T) -> R) -> R;
    }

    mod store {
        use std::{
            any::{Any, TypeId},
            cell::{Cell, RefCell},
            collections::HashMap,
        };

        use slab::Slab;

        use super::{Key, Owner};

        struct Item<T> {
            value: T,
            owners_count: Cell<usize>,
        }

        type TypedStore<T> = Slab<Item<T>>;

        thread_local!(
            static STORES: RefCell<
                HashMap<
                    TypeId,
                    // Box<TypedStore<T>>
                    Box<dyn Any>,
                >,
            > = Default::default();
        );

        fn store_mut<T: 'static, R>(f: impl FnOnce(&mut TypedStore<T>) -> R) -> R {
            STORES.with_borrow_mut(|stores| {
                let store = stores
                    .entry(TypeId::of::<T>())
                    .or_insert_with(|| Box::new(TypedStore::<T>::new()))
                    .downcast_mut::<TypedStore<T>>()
                    .unwrap();
                f(store)
            })
        }

        fn store<T: 'static, R>(f: impl FnOnce(&TypedStore<T>) -> R) -> R {
            STORES.with_borrow(|stores| {
                let store = stores
                    .get(&TypeId::of::<T>())
                    .unwrap()
                    .downcast_ref::<TypedStore<T>>()
                    .unwrap();
                f(store)
            })
        }

        impl super::StoreImpl for super::Store {
            fn insert<T: 'static>(value: T) -> Owner<T> {
                store_mut::<T, Owner<T>>(|s| {
                    let index = s.insert(Item {
                        value,
                        owners_count: Cell::new(1),
                    });
                    let key = Key::new(index);
                    Owner { key }
                })
            }

            fn unregister_owner<T>(owner: &mut Owner<T>) {
                let index = owner.key.index;

                let should_remove = store::<T, bool>(|items| {
                    let item = items.get(index).unwrap();

                    let new_count = item.owners_count.get() - 1;

                    if new_count == 0 {
                        true
                    } else {
                        item.owners_count.set(new_count);
                        false
                    }
                });

                if should_remove {
                    store_mut::<T, ()>(|items| {
                        let _: T = items.remove(index).value;
                    })
                }
            }

            fn register_owner_at<T>(key: Key<T>) -> Owner<T> {
                store::<T, ()>(|items| {
                    let item = items.get(key.index).unwrap();

                    item.owners_count.set(item.owners_count.get() + 1);
                });

                Owner { key }
            }

            fn owners_count_at<T: 'static>(key: Key<T>) -> usize {
                store::<T, usize>(|items| items.get(key.index).unwrap().owners_count.get())
            }

            fn map_at<T: 'static, R>(key: Key<T>, f: impl FnOnce(&T) -> R) -> R {
                store::<T, R>(|items| f(&items.get(key.index).unwrap().value))
            }

            fn map_mut_at<T: 'static, R>(key: Key<T>, f: impl FnOnce(&mut T) -> R) -> R {
                store_mut::<T, R>(|items| f(&mut items.get_mut(key.index).unwrap().value))
            }
        }
    }

    impl<T: 'static> super::Store<T> for Store {
        type Key = Key<T>;

        fn insert(&self, value: T) -> Owner<T> {
            <Store as StoreImpl>::insert::<T>(value)
        }
    }

    impl<T: 'static> super::Owner for Owner<T> {
        type Key = Key<T>;

        fn key(&self) -> Self::Key {
            self.key
        }

        fn shared_count(&self) -> usize {
            Store::owners_count_at(self.key)
        }
    }

    impl<T: 'static> Clone for Owner<T> {
        fn clone(&self) -> Self {
            use super::Key;
            self.key.owner()
        }
    }

    impl<T: 'static> Drop for Owner<T> {
        fn drop(&mut self) {
            Store::unregister_owner::<T>(self)
        }
    }
}