use std::fmt;
use std::hash::Hash;
use std::mem::ManuallyDrop;
use std::ops::Deref;
use std::ops::Index;

use merc_utilities::GenerationCounter;
use merc_utilities::GenerationalIndex;

/// A type-safe index for the ProtectionSet to prevent accidental use of wrong indices
#[repr(transparent)]
#[derive(Copy, Clone, Default, PartialEq, Eq, Hash)]
pub struct ProtectionIndex(GenerationalIndex<usize>);

impl Deref for ProtectionIndex {
    type Target = usize;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl fmt::Debug for ProtectionIndex {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "ProtectionIndex({:?})", self.0)
    }
}

impl fmt::Display for ProtectionIndex {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.0)
    }
}

/// A collection that assigns a unique index to every object added to it, and allows
/// removing objects while reusing their indices later. This is useful for managing
/// objects that must not be garbage collected, and as such it is called a protection set.
/// It is similar to a [`merc_collections::IndexedSet`], except that we cannot look up elements by value.
#[derive(Default)]
pub struct ProtectionSet<T> {
    roots: Vec<Entry<T>>, // The set of root active nodes.
    free: Option<usize>,
    number_of_insertions: u64,
    size: usize,

    /// The number of generations
    generation_counter: GenerationCounter,
}

impl<T> ProtectionSet<T> {
    /// Creates a new empty protection set.
    pub fn new() -> Self {
        ProtectionSet {
            roots: Vec::new(),
            free: None,
            number_of_insertions: 0,
            size: 0,
            generation_counter: GenerationCounter::new(),
        }
    }

    /// Returns the number of insertions into the protection set.
    pub fn number_of_insertions(&self) -> u64 {
        self.number_of_insertions
    }

    /// Returns maximum number of active instances.
    pub fn maximum_size(&self) -> usize {
        self.roots.capacity()
    }

    /// Returns the number of roots in the protection set
    pub fn len(&self) -> usize {
        self.size
    }

    /// Returns whether the protection set is empty.
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Adds the given object to the protection set and returns its index.
    pub fn protect(&mut self, object: T) -> ProtectionIndex {
        self.number_of_insertions += 1;
        self.size += 1;

        let index = match self.free {
            Some(first) => {
                let next = unsafe { self.roots[first].next };
                if first == next {
                    // The list is empty as its first element points to itself.
                    self.free = None;
                } else {
                    // Update free to be the next element in the list.
                    self.free = Some(next);
                }

                self.roots[first] = Entry {
                    object: ManuallyDrop::new(object),
                };
                first
            }
            None => {
                // If free list is empty insert new entry into roots.
                self.roots.push(Entry {
                    object: ManuallyDrop::new(object),
                });

                self.roots.len() - 1
            }
        };

        ProtectionIndex(self.generation_counter.create_index(index))
    }

    /// Remove protection from the given object. Note that index must be the
    /// index returned by the [ProtectionSet::protect] call.
    pub fn unprotect(&mut self, index: ProtectionIndex) {
        let index = self.generation_counter.get_index(index.0);
        self.size -= 1;

        // SAFETY: `index` refers to an occupied entry. We must drop the stored
        // object before overwriting the union slot with freelist metadata.
        unsafe {
            ManuallyDrop::drop(&mut self.roots[index].object);
        }

        match self.free {
            Some(next) => {
                self.roots[index] = Entry { next };
            }
            None => {
                self.roots[index] = Entry { next: index };
            }
        };

        self.free = Some(index);

        // Postcondition: verify the object was correctly removed from protection
        debug_assert!(
            self.freelist_iter().any(|free_idx| free_idx == index),
            "Failed to unprotect object"
        );
    }

    /// Replaces the object at the given index with the new object.
    pub fn replace(&mut self, index: ProtectionIndex, object: T) {
        let index = self.generation_counter.get_index(index.0);

        debug_assert!(
            !self.freelist_iter().any(|free_idx| free_idx == index),
            "Index {index} does not point to a filled entry"
        );

        self.roots[index] = Entry {
            object: ManuallyDrop::new(object),
        };
    }

    /// Returns an iterator over all root indices in the protection set.
    ///
    /// This is an O(n) operation in the size of the freelist, so it should be
    /// used with care. It is intended for debugging and testing purposes.
    pub fn iter(&self) -> ProtSetIter<'_, T> {
        let mut free_indices: Vec<usize> = self.freelist_iter().collect();
        free_indices.sort_unstable();
        ProtSetIter {
            current: 0,
            protection_set: self,
            generation_counter: &self.generation_counter,
            free_indices,
        }
    }

    /// Returns an iterator over all free indices in the freelist.
    ///
    /// This is intended for use in `debug_assert!`s to verify that an entry is
    /// filled (not in the freelist) or free (in the freelist).
    fn freelist_iter(&self) -> FreeListIter<'_, T> {
        FreeListIter {
            current: self.free,
            protection_set: self,
        }
    }

    /// Returns whether the protection set contains the given index.
    ///
    /// This operation is O(n) in the size of the freelist.
    pub fn contains_root(&self, index: ProtectionIndex) -> bool {
        let idx = self.generation_counter.get_index(index.0);
        !self.freelist_iter().any(|free_idx| free_idx == idx)
    }
}

impl<T> Drop for ProtectionSet<T> {
    fn drop(&mut self) {
        for index in 0..self.roots.len() {
            if self.freelist_iter().any(|free_idx| free_idx == index) {
                continue;
            }

            // SAFETY: indices not present in the freelist still contain a live
            // protected object that must be dropped.
            unsafe {
                ManuallyDrop::drop(&mut self.roots[index].object);
            }
        }
    }
}

impl<T> Index<ProtectionIndex> for ProtectionSet<T> {
    type Output = T;

    fn index(&self, index: ProtectionIndex) -> &Self::Output {
        let idx = *index;
        debug_assert!(
            !self.freelist_iter().any(|free_idx| free_idx == idx),
            "Attempting to index free spot {}",
            index,
        );
        // SAFETY: The generational index ensures this slot was not freed after the index was issued.
        unsafe { &self.roots[idx].object }
    }
}

impl<T: fmt::Debug> fmt::Debug for ProtectionSet<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_map().entries(self.iter()).finish()
    }
}

/// An entry in the protection set, which can either be filled with an object or
/// free and point to the next free entry.
union Entry<T> {
    object: ManuallyDrop<T>,

    // The next free entry in the free list. If this points to itself, the free list is empty.
    next: usize,
}

/// Check that the Entry is the same size as a usize.
#[cfg(not(debug_assertions))]
const _: () = assert!(std::mem::size_of::<Entry<usize>>() == std::mem::size_of::<usize>());

/// An iterator over the filled entries in a protection set.
pub struct ProtSetIter<'a, T> {
    current: usize,
    protection_set: &'a ProtectionSet<T>,
    generation_counter: &'a GenerationCounter,
    /// Sorted free indices collected at iterator construction to skip free slots.
    free_indices: Vec<usize>,
}

impl<'a, T> Iterator for ProtSetIter<'a, T> {
    type Item = (ProtectionIndex, &'a T);

    fn next(&mut self) -> Option<Self::Item> {
        // Find the next valid entry, return it when found or None when end of roots is reached.
        while self.current < self.protection_set.roots.len() {
            let idx = self.current;
            self.current += 1;

            if self.free_indices.binary_search(&idx).is_err() {
                // SAFETY: idx is not in the free list, so it holds a valid filled object.
                let object = unsafe { &*self.protection_set.roots[idx].object };
                return Some((ProtectionIndex(self.generation_counter.recall_index(idx)), object));
            }
        }

        None
    }
}

/// An iterator over the free indices in the protection set freelist.
struct FreeListIter<'a, T> {
    current: Option<usize>,
    protection_set: &'a ProtectionSet<T>,
}

impl<'a, T> Iterator for FreeListIter<'a, T> {
    type Item = usize;

    fn next(&mut self) -> Option<Self::Item> {
        let curr = self.current?;
        // SAFETY: curr is a valid index in the free list, so it stores a `next` pointer.
        let next = unsafe { self.protection_set.roots[curr].next };
        if next == curr {
            // Sentinel: last free entry points to itself.
            self.current = None;
        } else {
            self.current = Some(next);
        }
        Some(curr)
    }
}

impl<'a, T> IntoIterator for &'a ProtectionSet<T> {
    type Item = (ProtectionIndex, &'a T);
    type IntoIter = ProtSetIter<'a, T>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

#[cfg(test)]
mod tests {
    use rand::RngExt;

    use merc_utilities::random_test;
    use merc_utilities::test_logger;

    use crate::ProtectionIndex;
    use crate::ProtectionSet;

    #[test]
    #[cfg_attr(miri, ignore)]
    fn test_random_protection_set() {
        random_test(100, |rng| {
            let mut protection_set = ProtectionSet::<usize>::new();

            // Protect a number of indices and record their roots.
            let mut indices: Vec<ProtectionIndex> = Vec::new();

            for _ in 0..5000 {
                indices.push(protection_set.protect(rng.random_range(0..1000)));
            }

            // Unprotect a number of roots.
            for index in 0..2500 {
                assert!(protection_set[indices[index]] <= 1000);
                protection_set.unprotect(indices[index]);
                indices.remove(index);
            }

            // Protect more to test the freelist
            for _ in 0..1000 {
                indices.push(protection_set.protect(rng.random_range(0..1000)));
            }

            for index in &indices {
                assert!(
                    protection_set.contains_root(*index),
                    "All indices that are not unprotected should occur in the protection set"
                );
            }

            assert_eq!(
                protection_set.iter().count(),
                6000 - 2500,
                "This is the number of roots remaining"
            );
            assert_eq!(protection_set.number_of_insertions(), 6000);
            assert!(protection_set.maximum_size() >= 5000);
            assert!(!protection_set.is_empty());
        });
    }

    #[test]
    fn test_protection_set_basic() {
        test_logger();

        let mut set = ProtectionSet::<String>::new();

        // Protect some values
        let idx1 = set.protect(String::from("value1"));
        let idx2 = set.protect(String::from("value2"));

        // Verify contains_root works
        assert!(set.contains_root(idx1));
        assert!(set.contains_root(idx2));

        // Test indexing
        assert_eq!(set[idx1], "value1");
        assert_eq!(set[idx2], "value2");

        // Test unprotect
        set.unprotect(idx1);
        assert!(!set.contains_root(idx1));
        assert!(set.contains_root(idx2));

        // Re-use freed slot
        let idx3 = set.protect(String::from("value3"));
        assert_eq!(set[idx3], "value3");
    }
}
