use std::marker::PhantomData;
use alloc::raw_vec::RawVec;

/// Entry in a freelist. Contains a generation number to ensure silly mistakes
/// using an item's index after freeing it.
struct Entry<T> {
    generation: usize,
    item: T,
}

/// Key of an item in the freelist. Contains a generation number to prevent
/// use-after-free during development.
// TODO: Two usizes are probably overkill
#[derive(Copy, Clone)]
pub struct Key<T> {
    generation: usize,
    index: usize,
    _type: PhantomData<T>,
}

/// Generic freelist structure. Item insertion/retrieval/deletion works like a
/// HashMap through keys.
pub struct FreeList<T> {
    items: *mut Entry<T>,
    capacity: usize,
    length: usize,
    free: Vec<usize>,
}

impl<T> FreeList<T> {
    pub fn new() -> Self<T> {
        std::alloc::Layout::from_size_align()
        Self{
            items: std::ptr::null_mut(),
            capacity: 0,
            length: 0,
            free: Vec::new(),
        }
    }

    pub fn with_capacity(capacity: usize) -> Self {
        alloc::
        Self{
            items: std::,
            free: Vec::with_capacity(capacity),
            length: 0,
        }
    }

    /// Inserts a new item into the freelist. Will return a key that can be used
    /// to retrieve the item and delete it.
    pub fn insert(&mut self, item: T) -> Key<T> {
        let mut generation;
        let mut index;

        if self.free.is_empty() {
            // No free elements, make sure we have enough capacity
            if self.length == self.items.capacity() {
                self.items.reserve(self.length, 1);
            }

            // Now we do
            generation = 0;
            index = self.length;

            unsafe {
                let target = self.items.ptr().add(self.length);
                std::ptr::write(&mut target.item, item);
                self.length += 1;
            }
        } else {
            // We have a free spot. Note that the generation is incremented upon
            // freeing an item. So we can just take the current generation value
            // here.
            index = self.free.pop().unwrap();

            unsafe {
                let target = self.items.ptr().add(self.length);
                generation = target.generation;
                std::ptr::write(&mut target.item, item);
            }
        }

        Key { generation, index, _type: PhantomData::default() }
    }

    /// Removes the entry using the provided key. Will panic if the element was
    /// removed already.
    pub fn erase(&mut self, index: Key<T>) {
        // This should always be the case
        debug_assert!(index.index < self.length);

        // Retrieve element and make sure that the generation matches
        unsafe {
            let entry = self.items.ptr().add(index.index);
            assert_eq!(entry.generation, entry.generation);
            *entry.generation += 1;
            std::ptr::drop_in_place(&mut entry.item);
        }

        // Add the entry to the freelist.
        self.free.push(index.index);
    }
}

impl<T> std::ops::Index<Key<T>> for FreeList<T> {
    type Output = T;

    fn index(&self, index: &Key<T>) -> &Self::Output {
        debug_assert!(index.index < self.length);
        unsafe {
            let entry = self.items.ptr().add(index.index);
            assert_eq!(entry.generation, index.generation);
            return &entry.item;
        }
    }
}

impl<T> std::ops::IndexMut<Key<T>> for FreeList<T> {
    fn index_mut(&self, index: &Key<T>) -> &Self::Output {
        debug_assert!(index.index < self.length);
        unsafe {
            let entry = self.items.ptr().add(index.index);
            assert_eq!(entry.generation, index.generation);
            return &mut entry.item;
        }
    }
}

impl<T> Drop for FreeList<T> {
    fn drop(&mut self) {
        // Sort free indices to use them while traversing the allocated items
        self.free.sort_unstable();
        let free_length = self.free.len();
        let mut next_free_idx = 1;
        let mut next_free_item = usize::MAX;
        if free_length != 0 {
            next_free_item = self.free[0];
        }

        // Go through all items. If we didn't yet drop the item, then we do
        // so here.
        for item_idx in 0..self.length {
            if item_idx == next_free_item {
                if next_free_idx < free_length {
                    next_free_item = self.free[next_free_idx];
                    next_free_idx += 1;
                } else {
                    next_free_item = usize::MAX;
                }

                // Skipped the current item, go to the next one
                continue
            }

            // Need to deallocate the current item
            unsafe {
                let entry = self.items.ptr().add(item_idx);
                std::ptr::drop_in_place(&mut entry.item);
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Arc;
    use std::sync::atomic::{AtomicU32, Ordering};

    struct Counters {
        constructed: Arc<AtomicU32>,
        destructed: Arc<AtomicU32>,
    }

    impl Counters {
        pub fn new() -> Self {
            Self{
                constructed: Arc::new(AtomicU32::new(0)),
                destructed: Arc::new(AtomicU32::new(0)),
            }
        }
    }

    struct TestEntity {
        counters: Counters,
        pub value: u32,
    }

    impl TestEntity {
        pub fn new(counters: &Counters, value: u32) -> Self {
            counters.constructed.fetch_add(1, Ordering::SeqCst);
            return TestEntity{
                counters: Counters{
                    constructed: counters.constructed.clone(),
                    destructed: counters.destructed.clone(),
                },
                value,
            }
        }
    }

    impl Drop for TestEntity {
        fn drop(&mut self) {
            self.counters.destructed.fetch_add(1, Ordering::SeqCst);
        }
    }

    #[test]
    fn only_constructing() {
        const NUM_CREATED: u32 = 3;
        let counters = Counters::new();
        let mut list = FreeList::new();

        for i in 0..NUM_CREATED {
            list.insert(TestEntity::new(&counters, i));
            assert_eq!(counters.constructed.load(Ordering::AcqRel), i + 1);
        }

        // Everything is constructed, check freelist
        assert_eq!(counters.constructed.load(Ordering::AcqRel), NUM_CREATED);
        assert_eq!(list.length, 3);
        assert!(list.items.capacity() >= 3);
        assert!(list.free.is_empty());

        // Drop and check everything is properly dropped
        drop(list);
        assert_eq!(counters.destructed.load(Ordering::AcqRel), NUM_CREATED)
    }

    #[test]
    fn reusing_slots() {
        const NUM_ROUNDS: u32 = 10;
        const NUM_IN_USE: u32 = 10;

        let counters = Counters::new();
        let mut list = FreeList::new();
        let mut indices = Vec::with_capacity(NUM_IN_USE as usize);

        for round_idx in 0..NUM_ROUNDS {
            indices.clear();

            // Adding entries
            for i in 0..NUM_IN_USE {
                let new_index = list.insert(TestEntity::new(&counters, i));
                indices.push(new_index);
            }

            // Length should always remain the same as the total number of
            // entries in use
            assert_eq!(list.length, NUM_IN_USE as usize);

            // Removing entries, and making sure that everything is still
            // accessible
            for idx in 0..NUM_IN_USE {
                // Make sure we can still retrieve the item we're going to delete
                let idx_to_remove = NUM_IN_USE - 1 - idx;
                let pop_index = indices.pop().unwrap();
                let entry = &list[pop_index];
                assert_eq!(entry.value, idx_to_remove);

                // Remove the entry and make sure the other ones are still
                // accessible
                list.erase(pop_index);

                for remaining_idx in 0..idx_to_remove + 1 {
                    let remaining_key = &indices[remaining_idx as usize];
                    let remaining_entry = &list[*remaining_key];
                    assert_eq!(remaining_entry.value, remaining_idx);
                }
            }

            // Now that we're empty, our constructed and destructed counts
            // should match
            let expected_count = (round_idx + 1) * NUM_IN_USE;
            assert_eq!(counters.constructed.load(Ordering::AcqRel), expected_count);
            assert_eq!(counters.destructed.load(Ordering::AcqRel), expected_count);
        }

        // Make sure the capacity didn't grow out of bounds, maximum growth rate
        // I've ever encountered on `Vec`s is 2. So:
        assert!(list.items.capacity() >= NUM_IN_USE as usize && list.items.capacity() < 2*NUM_IN_USE as usize);

        // Finally, when we drop the list we shouldn't be destructing anything
        // anymore.
        drop(list);
        let final_count = NUM_ROUNDS * NUM_IN_USE;
        assert_eq!(counters.constructed.load(Ordering::AcqRel), final_count);
        assert_eq!(counters.destructed.load(Ordering::AcqRel), final_count);
    }

    #[test]
    #[should_panic]
    fn panic_on_reused_key_of_empty_slot() {
        let counters = Counters::new();
        let mut list = FreeList::new();
        let key = list.insert(TestEntity::new(&counters, 0));
        list.erase(key);
        let entry = &list[key];
    }

    #[test]
    #[should_panic]
    fn panic_on_reused_key_of_used_slot() {
        let counters = Counters::new();
        let mut list = FreeList::new();
        let key1 = list.insert(TestEntity::new(&counters, 0));
        list.erase(key1);
        let key2 = list.insert(TestEntity::new(&counters, 0));
        assert_eq!(key1.index, key2.index);
        assert_ne!(key1.generation, key2.generation);
        let entry = &list[key1];
    }
}