pub trait AdventIterator: Iterator {
    fn into_chunks<const N: usize>(self) -> impl Iterator<Item = Chunck<N, Self::Item>>
    where
        Self: Sized,
    {
        in_chunks::<N, Self::Item>(self)
    }
}

impl<I> AdventIterator for I where I: Iterator {}

#[derive(Debug, PartialEq, Eq)]
pub enum Chunck<const N: usize, T> {
    Next([T; N]),
    Rest(Vec<T>),
}

pub fn in_chunks<const N: usize, T>(
    mut iterator: impl Iterator<Item = T>,
) -> impl Iterator<Item = Chunck<N, T>> {
    let mut buffer = Vec::new();
    let mut done = false;
    std::iter::from_fn(move || {
        if done {
            return None;
        }
        for _ in 0..N {
            match iterator.next() {
                Some(to_push) => buffer.push(to_push),
                None => {
                    done = true;
                    return if buffer.is_empty() {
                        None
                    } else {
                        Some(Chunck::Rest(std::mem::take(&mut buffer)))
                    };
                }
            }
        }
        let array: [T; N] = std::mem::take(&mut buffer)
            .try_into()
            .unwrap_or_else(|_| panic!("Buffer must have the same size as the N ({})", N));
        Some(Chunck::Next(array))
    })
}

#[cfg(test)]
mod testing {
    use crate::iterations::{in_chunks, Chunck};

    #[test]
    fn should_split_chunks() {
        fn assert_case<const N: usize>(input: Vec<u32>, expected: Vec<Chunck<N, u32>>) {
            let actual: Vec<Chunck<N, u32>> = in_chunks::<N, u32>(input.into_iter()).collect();
            assert_eq!(expected, actual);
        }

        assert_case::<2>(vec![2], vec![Chunck::Rest(vec![2])]);
        assert_case::<2>(vec![], vec![]);
        assert_case::<2>(vec![2, 2], vec![Chunck::Next([2, 2])]);
        assert_case::<2>(
            vec![1, 2, 3],
            vec![Chunck::Next([1, 2]), Chunck::Rest(vec![3])],
        );
        assert_case::<1>(
            vec![1, 2, 3],
            vec![Chunck::Next([1]), Chunck::Next([2]), Chunck::Next([3])],
        );
    }
}