use super::Entity;

pub struct SparseSet<T> {
    pub dense: Vec<T>,
    pub entities: Vec<Entity>,
    pub sparse: Vec<Option<usize>>,
}

impl<T> SparseSet<T> {
    pub fn new() -> Self {
        Self {
            dense: Vec::new(),
            entities: Vec::new(),
            sparse: Vec::new(),
        }
    }

    pub fn insert(&mut self, entity: Entity, component: T) {
        if self.sparse.len() <= entity {
            self.sparse.resize(entity + 1, None);
        }

        self.dense.push(component);
        self.entities.push(entity);

        let id = self.dense.len() - 1;
        self.sparse[entity] = Some(id);
    }

    pub fn remove(&mut self, entity: Entity) {
        let Some(id) = self.sparse.get(entity).copied().flatten() else {
            return;
        };
        let Some(last) = self.entities.last().copied() else {
            return;
        };
        self.dense.swap_remove(id);
        self.entities.swap_remove(id);
        self.sparse[last] = Some(id);
        self.sparse[entity] = None;
    }

    pub fn iter(&self) -> impl Iterator<Item = (Entity, &T)> {
        self.entities.iter().copied().zip(self.dense.iter())
    }

    pub fn iter_mut(&mut self) -> impl Iterator<Item = (Entity, &mut T)> {
        self.entities.iter().copied().zip(self.dense.iter_mut())
    }

    pub fn get(&self, entity: Entity) -> Option<&T> {
        if let Some(id) = self.sparse.get(entity) {
            match id {
                Some(id) => self.dense.get(*id),
                None => None,
            }
        } else {
            None
        }
    }

    pub fn get_mut(&mut self, entity: Entity) -> Option<&mut T> {
        if let Some(id) = self.sparse.get(entity) {
            match id {
                Some(id) => self.dense.get_mut(*id),
                None => None,
            }
        } else {
            None
        }
    }
}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn test_remove() {
        let mut set = SparseSet::<i32>::new();
        set.insert(3, 300);
        set.insert(2, 200);
        set.insert(4, 400);
        set.insert(7, 700);

        set.remove(4);

        assert!(set.get(7).copied() == Some(700));
        assert!(set.get(3).copied() == Some(300));
        assert!(set.get(2).copied() == Some(200));
        assert!(set.get(4).copied().is_none());

        set.remove(7);
        assert!(set.get(7).copied().is_none());
        assert!(set.get(3).copied() == Some(300));
        assert!(set.get(2).copied() == Some(200));
        assert!(set.get(4).copied().is_none());
    }
}