use core::slice;

use argon2::Block;

pub(crate) struct Argon2Blocks<A: alloc::alloc::Allocator> {
    p: core::ptr::NonNull<Block>,
    len: usize,
    alloc: A,
}

impl<A: alloc::alloc::Allocator> Argon2Blocks<A> {
    /// Each block is 1 KiB.
    /// Total size is `len` * 1 KiB
    pub fn new_in(len: usize, alloc: A) -> Option<Self> {
        use alloc::alloc::Layout;

        if len == 0 {
            return None;
        }

        let layout = Layout::array::<Block>(len).ok()?;
        // SAFETY: `alloc_zeroed` is used correctly with non-zero layout
        let p = alloc.allocate_zeroed(layout).ok()?.cast();

        Some(Self { p, len, alloc })
    }
}

impl<A: alloc::alloc::Allocator> AsMut<[Block]> for Argon2Blocks<A> {
    fn as_mut(&mut self) -> &mut [Block] {
        unsafe { slice::from_raw_parts_mut(self.p.as_ptr(), self.len) }
    }
}

impl<A: alloc::alloc::Allocator> AsRef<[Block]> for Argon2Blocks<A> {
    fn as_ref(&self) -> &[Block] {
        unsafe { slice::from_raw_parts(self.p.as_ptr(), self.len) }
    }
}

impl<A: alloc::alloc::Allocator> Drop for Argon2Blocks<A> {
    fn drop(&mut self) {
        use alloc::alloc::Layout;
        // SAFETY: layout was checked during construction
        let layout = unsafe { Layout::array::<Block>(self.len).unwrap_unchecked() };
        // SAFETY: we use `dealloc` correctly with the previously allocated pointer
        unsafe {
            self.alloc.deallocate(self.p.cast(), layout);
        }
    }
}
