use crate::node::Cells;

#[derive(Debug, Clone, Copy)]
/// An iterator that yields `N` different-sized [`Cells`] from a single [`Cells`] buffer, and a given cell count for each [`Cells`] slice.
pub(crate) struct NCellsIter<'a, const N: usize> {
    cells: Cells<'a>,
    cell_count: [u32; N],
}

impl<'a, const N: usize> NCellsIter<'a, N> {
    pub const fn new(cells: Cells<'a>, cell_count: [u32; N]) -> Self {
        Self { cells, cell_count }
    }
}

impl<'a, const N: usize> Iterator for NCellsIter<'a, N> {
    type Item = [Cells<'a>; N];

    fn next(&mut self) -> Option<Self::Item> {
        let mut result = [const { Cells::new(&[]) }; N];
        for i in 0..N {
            let (stole, remaining) = self.cells.try_spilt_at(self.cell_count[i] as usize)?;
            result[i] = stole;
            self.cells = remaining;
        }
        Some(result)
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub(crate) struct RangesIter<'a> {
    cells: Cells<'a>,
    bus_address_cells: u32,
    address_cells: u32,
    bus_size_cells: u32,
}

impl<'a> RangesIter<'a> {
    pub const fn new(
        cells: Cells<'a>,
        bus_address_cells: u32,
        address_cells: u32,
        bus_size_cells: u32,
    ) -> Option<Self> {
        if address_cells as usize > (size_of::<usize>() / size_of::<u32>())
            || bus_size_cells as usize > (size_of::<usize>() / size_of::<u32>())
        {
            return None;
        }
        Some(Self {
            cells,
            bus_address_cells,
            address_cells,
            bus_size_cells,
        })
    }
}

impl<'a> Iterator for RangesIter<'a> {
    type Item = (Cells<'a>, usize, usize);

    fn next(&mut self) -> Option<Self::Item> {
        if (self.bus_address_cells + self.address_cells + self.bus_size_cells) as usize
            > self.cells.len()
        {
            return None;
        }

        let bus_address_cells = self.cells.take_n(self.bus_address_cells)?;
        let address = self.cells.next_usize(self.address_cells)?;
        let size = self.cells.next_usize(self.bus_size_cells)?;

        Some((bus_address_cells, address, size))
    }
}
