use std::ops::{Index, IndexMut};

/// An AST node ID for lookups in the arena
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct NodeId(pub usize);

/// Arenas are efficient ways of storing trees: All nodes are actually on a
/// flat surface, but have IDs to other nodes.
#[derive(Debug)]
pub enum Arena<'a, T: 'a> {
    Owned(Vec<Option<T>>),
    Borrowed(&'a mut Vec<Option<T>>)
}
impl<'a, T> Default for Arena<'a, T> {
    fn default() -> Self {
        Arena::Owned(Vec::new())
    }
}
impl<'a, T> Arena<'a, T> {
    /// Create a new arena instance
    pub fn new() -> Self {
        Self::default()
    }
    /// Create a new arena instance that reuses an existing vector
    pub fn with_arena(arena: &'a mut Vec<Option<T>>) -> Self {
        Arena::Borrowed(arena)
    }
    /// Create a new arena reference that shares the inner vector. This
    /// reference will make the existing instance invalid for use while the
    /// reference is active.
    pub fn reference<'b>(&'b mut self) -> Arena<'b, T> {
        Arena::with_arena(self.get_mut())
    }

    /// Return the arena instance
    ///
    /// # Panics
    /// Panics if this arena instance was created with Borrowed
    pub fn into_inner(self) -> Vec<Option<T>> {
        match self {
            Arena::Owned(inner) => inner,
            Arena::Borrowed(_) => panic!("can't move out of borrowed arena")
        }
    }
    /// Return a slice pointing to the inner arena representation
    pub fn get_ref(&self) -> &[Option<T>] {
        match self {
            Arena::Owned(inner) => inner,
            Arena::Borrowed(inner) => inner
        }
    }
    fn get_mut(&mut self) -> &mut Vec<Option<T>> {
        match self {
            Arena::Owned(inner) => inner,
            Arena::Borrowed(inner) => inner
        }
    }
    /// Return an iterator over the nodes in this arena
    pub fn iter(&self) -> ArenaIter<T> {
        self.into_iter()
    }
    /// Place an element into the arena
    pub fn insert(&mut self, elem: T) -> NodeId {
        let inner = self.get_mut();
        inner.push(Some(elem));
        NodeId(inner.len() - 1)
    }
    /// Move out of an element in the arena. This will make following indexing
    /// calls on that index panic. This will not deallocate any space. This is
    /// due to the fact that all indexes need to stay the same. We could use a
    /// different data type instead of a tree to prevent this, but that instead
    /// means lookup will be slightly slower.
    pub fn take(&mut self, index: NodeId) -> T {
        self.get_mut()[index.0].take().expect("this index has been taken")
    }

    /// Get an iterator over all existing NodeIds in the arena.
    pub fn ids(&'a self) -> impl Iterator<Item = NodeId> + 'a {
        let inner = self.get_ref();

        inner
            .iter()
            .enumerate()
            .filter(|(_, val)| val.is_some())
            .map(|(id, _)| NodeId(id))
    }
}
impl<'a, T> Index<NodeId> for Arena<'a, T> {
    type Output = T;

    fn index(&self, index: NodeId) -> &Self::Output {
        &self.get_ref()[index.0].as_ref().expect("this index has been taken")
    }
}
impl<'a, T> IndexMut<NodeId> for Arena<'a, T> {
    fn index_mut(&mut self, index: NodeId) -> &mut Self::Output {
        self.get_mut()[index.0].as_mut().expect("this index has been taken")
    }
}
impl<'r, 'a, T> IntoIterator for &'r Arena<'a, T> {
    type Item = &'r T;
    type IntoIter = ArenaIter<'r, T>;

    fn into_iter(self) -> Self::IntoIter {
        ArenaIter {
            slice: self.get_ref(),
            index: 0
        }
    }
}

impl<'a, T> Clone for Arena<'a, T> where T: Clone {
    fn clone(&self) -> Self {
        match self {
            Arena::Owned(inner) => Arena::Owned(inner.to_vec()),
            Arena::Borrowed(inner) => Arena::Owned(inner.to_vec())
        }
    }
}

/// An iterator over an arena, created with `iter`
pub struct ArenaIter<'a, T: 'a> {
    slice: &'a [Option<T>],
    index: usize
}
impl<'a, T> Iterator for ArenaIter<'a, T> {
    type Item = &'a T;

    fn next(&mut self) -> Option<Self::Item> {
        while let Some(None) = self.slice.get(self.index) {
            self.index += 1;
        }
        let item = self.slice.get(self.index).map(|item| item.as_ref().unwrap());
        self.index += 1;
        item
    }
}

#[cfg(test)]
mod tests {
    use super::{Arena, NodeId};

    #[test]
    fn it_should_return_node_ids() {
        let mut arena: Arena<'static, u32> = Arena::new();

        arena.insert(100);
        arena.insert(101);

        let ids: Vec<NodeId> = arena.ids().collect();
        assert_eq!(
            ids,
            vec![NodeId(0), NodeId(1)]
        )
    }

    #[test]
    fn it_should_not_return_node_ids_of_removed_elements() {
        let mut arena: Arena<'static, u32> = Arena::new();

        arena.insert(100);
        arena.insert(101);
        arena.take(NodeId(0));

        let ids: Vec<NodeId> = arena.ids().collect();
        assert_eq!(
            ids,
            vec![NodeId(1)]
        )
    }
}