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

enum Cell<T> {
    Free { next_free_cell_id: Option<usize> },
    Value { value: T },
}

pub struct _PreserveIndexVec<T> {
    next_free: Option<usize>,
    buffer: Vec<Cell<T>>,
    len: usize,
}

impl<T> _PreserveIndexVec<T> {
    pub fn new() -> Self {
        return Self {
            next_free: None,
            buffer: Vec::new(),
            len: 0,
        };
    }
    pub fn len(&self) -> usize {
        return self.len;
    }

    pub fn add(&mut self, constructor: impl FnOnce(&mut Self, usize) -> T) -> usize {
        match self.next_free {
            None => {
                let cell = Cell::Value {
                    value: constructor(self, self.len),
                };
                self.buffer.push(cell);
                self.len += 1;
                return self.buffer.len() - 1;
            }
            Some(i) => match self.buffer[i] {
                Cell::Free { next_free_cell_id } => {
                    self.next_free = next_free_cell_id;
                    let cell = Cell::Value {
                        value: constructor(self, i),
                    };
                    self.buffer[i] = cell;
                    self.len += 1;
                    return i;
                }
                Cell::Value { .. } => panic!("Busy cell marked as free"),
            },
        }
    }

    pub fn remove(&mut self, cell_index: usize) {
        match self.buffer[cell_index] {
            Cell::Free { .. } => panic!("Cell already free"),
            Cell::Value { .. } => {
                self.buffer[cell_index] = Cell::Free {
                    next_free_cell_id: self.next_free,
                };
                self.next_free = Some(cell_index);
                self.len -= 1;
            }
        }
    }

    pub fn iter(&self) -> impl Iterator<Item = &T> {
        return self.buffer.iter().filter_map(|c| match c {
            Cell::Free { .. } => return None,
            Cell::Value { value } => return Some(value),
        });
    }

    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut T> {
        return self.buffer.iter_mut().filter_map(|c| match c {
            Cell::Free { .. } => return None,
            Cell::Value { value } => return Some(value),
        });
    }

    pub fn iter_indexed(&self) -> impl Iterator<Item = (usize, &T)> {
        return self.buffer.iter().enumerate().filter_map(|(i, c)| match c {
            Cell::Free { .. } => return None,
            Cell::Value { value } => return Some((i, value)),
        });
    }

    pub fn iter_indexed_mut(&mut self) -> impl Iterator<Item = (usize, &mut T)> {
        return self
            .buffer
            .iter_mut()
            .enumerate()
            .filter_map(|(i, c)| match c {
                Cell::Free { .. } => return None,
                Cell::Value { value } => return Some((i, value)),
            });
    }

    pub fn iter_indexes(&self) -> impl Iterator<Item = (usize)> {
        return self.buffer.iter().enumerate().filter_map(|(i, c)| match c {
            Cell::Free { .. } => return None,
            Cell::Value { value } => return Some(i),
        });
    }

    pub fn capacity(&self) -> usize {
        return self.buffer.len();
    }
}

impl<T> Index<usize> for _PreserveIndexVec<T> {
    type Output = T;

    fn index(&self, index: usize) -> &Self::Output {
        match self.buffer.get(index).expect("Cell index out of bounds") {
            Cell::Free { .. } => panic!("Referenced cell doesn't exists"),
            Cell::Value { value } => return value,
        }
    }
}

impl<T> IndexMut<usize> for _PreserveIndexVec<T> {
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        match self
            .buffer
            .get_mut(index)
            .expect("Cell index out of bounds")
        {
            Cell::Free { .. } => panic!("Referenced cell doesn't exists"),
            Cell::Value { value } => return value,
        }
    }
}