//! Block management interfaces and structures
use core::{
    marker::PhantomData,
    mem::{size_of, MaybeUninit},
    ops::{Deref, DerefMut},
    ptr::NonNull,
};

use yggdrasil_abi::error::Error;

/// Number of bytes in a block
pub const SIZE: usize = 4096;
/// Maximum number of indirection pointers a block can hold
pub const ENTRY_COUNT: usize = SIZE / size_of::<usize>();

/// Interface for a block allocator
///
/// # Safety
///
/// This trait is unsafe to implement because it has to provide and accept raw data pointers of
/// exactly [SIZE].
pub unsafe trait BlockAllocator: Send + Sync + 'static {
    /// Allocates a contiguous block of size [SIZE]
    fn alloc() -> Result<NonNull<u8>, Error>;

    /// Dealocates a block.
    ///
    /// # Safety
    ///
    /// Unsafe: accepts arbitrary data pointers.
    unsafe fn dealloc(block: NonNull<u8>);
}

struct BlockRef<'a, A: BlockAllocator> {
    ptr: usize,
    _pd: PhantomData<&'a A>,
}

#[repr(transparent)]
struct BlockRaw<'a, A: BlockAllocator> {
    inner: BlockRef<'a, A>, // inner: Option<&'a mut [u8; SIZE]>,
}

/// Block containing file data
#[repr(transparent)]
pub struct BlockData<'a, A: BlockAllocator> {
    inner: BlockRaw<'a, A>,
}

/// Block containing indirection pointers to other blocks
#[repr(transparent)]
pub struct BlockIndirect<'a, A: BlockAllocator> {
    inner: BlockRaw<'a, A>,
}

/// Block containing indirection pointers to indirection pointers to data blocks
#[repr(transparent)]
pub struct BlockIndirectIndirect<'a, A: BlockAllocator> {
    inner: BlockRaw<'a, A>,
}

impl<'a, A: BlockAllocator> BlockRef<'a, A> {
    const fn null() -> Self {
        Self {
            ptr: 0,
            _pd: PhantomData,
        }
    }

    unsafe fn from_allocated(address: usize) -> Self {
        debug_assert_eq!(address & 1, 0);
        Self {
            ptr: address,
            _pd: PhantomData,
        }
    }

    unsafe fn copy_on_write(address: usize) -> Self {
        debug_assert_eq!(address & 1, 0);
        Self {
            ptr: address | 1,
            _pd: PhantomData,
        }
    }

    #[inline]
    fn is_allocated(&self) -> bool {
        self.ptr & 1 == 0
    }

    #[inline]
    fn is_null(&self) -> bool {
        self.ptr == 0
    }

    #[inline]
    fn as_mut(&mut self) -> &'a mut [u8; SIZE] {
        if self.is_null() {
            panic!("Null block dereference");
        }

        // FIXME: if a non-full block has been marked as CoW, the function will overrun the file
        // boundary
        if !self.is_allocated() {
            // Allocate the block
            let ptr = A::alloc().expect("Could not allocate a block").as_ptr() as usize;
            // Clone data
            let src = self.as_ref();
            let dst = unsafe { core::slice::from_raw_parts_mut(ptr as *mut u8, SIZE) };

            dst.copy_from_slice(src);

            self.ptr = ptr;
        }

        unsafe { &mut *((self.ptr & !1) as *mut [u8; SIZE]) }
    }

    #[inline]
    fn as_ref(&self) -> &'a [u8; SIZE] {
        if self.is_null() {
            panic!("Null block dereference");
        }

        unsafe { &*((self.ptr & !1) as *const [u8; SIZE]) }
    }
}

impl<A: BlockAllocator> Drop for BlockRef<'_, A> {
    fn drop(&mut self) {
        if self.is_allocated() && !self.is_null() {
            unsafe {
                A::dealloc(NonNull::new_unchecked(self.ptr as *mut _));
            }
        }
    }
}

impl<'a, A: BlockAllocator> BlockRaw<'a, A> {
    const fn null() -> Self {
        Self {
            inner: BlockRef::null(),
        }
    }

    fn new() -> Result<Self, Error> {
        let ptr = A::alloc()?;
        unsafe {
            Ok(Self {
                inner: BlockRef::from_allocated(ptr.as_ptr() as _),
            })
        }
    }

    unsafe fn as_uninit_indirect_mut(
        &mut self,
    ) -> &'a mut [MaybeUninit<BlockData<'a, A>>; ENTRY_COUNT] {
        if self.inner.is_null() {
            panic!("Null block dereference");
        }

        &mut *(self.inner.ptr as *mut _)
    }

    #[inline]
    unsafe fn as_data_ref(&self) -> &'a [u8; SIZE] {
        self.inner.as_ref()
    }

    #[inline]
    unsafe fn as_data_mut(&mut self) -> &'a mut [u8; SIZE] {
        self.inner.as_mut()
    }

    #[inline]
    fn is_null(&self) -> bool {
        self.inner.is_null()
    }
}

// Data block
impl<A: BlockAllocator> BlockData<'_, A> {
    /// Dummy entry representing a missing block
    pub const fn null() -> Self {
        Self {
            inner: BlockRaw::null(),
        }
    }

    /// Create a Copy-on-Write data block from existing data.
    ///
    /// # Safety
    ///
    /// This function is unsafe as it accepts arbitrary pointers. The caller must ensure the
    /// address is properly aligned (at least to a u16 boundary), does not cross any device memory
    /// and the pointer outlives the block reference.
    pub unsafe fn copy_on_write(address: usize) -> Self {
        Self {
            inner: BlockRaw {
                inner: BlockRef::copy_on_write(address),
            },
        }
    }

    /// Allocates a new block for data
    pub fn new() -> Result<Self, Error> {
        Ok(Self {
            inner: BlockRaw::new()?,
        })
    }

    /// Replaces self with a null block and drops any data that might've been allocated
    pub fn set_null(&mut self) {
        self.inner = BlockRaw::null();
    }

    /// Returns `true` if the block this structure refers to has not yet been allocated
    #[inline]
    pub fn is_null(&self) -> bool {
        self.inner.is_null()
    }
}

impl<A: BlockAllocator> Deref for BlockData<'_, A> {
    type Target = [u8; SIZE];

    fn deref(&self) -> &Self::Target {
        unsafe { self.inner.as_data_ref() }
    }
}

impl<A: BlockAllocator> DerefMut for BlockData<'_, A> {
    #[inline]
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { self.inner.as_data_mut() }
    }
}

// Indirect block
impl<A: BlockAllocator> BlockIndirect<'_, A> {
    /// Dummy entry representing a missing block
    pub const fn null() -> Self {
        Self {
            inner: BlockRaw::null(),
        }
    }

    /// Allocates a new indirection block
    pub fn new() -> Result<Self, Error> {
        let mut inner = BlockRaw::new()?;
        for item in unsafe { inner.as_uninit_indirect_mut() } {
            item.write(BlockData::null());
        }
        Ok(Self { inner })
    }

    /// Replaces self with a null block and drops any data that might've been allocated
    pub fn set_null(&mut self) {
        self.inner = BlockRaw::null();
    }

    /// Returns `true` if the block this structure refers to has not yet been allocated
    #[inline]
    pub fn is_null(&self) -> bool {
        self.inner.is_null()
    }
}

impl<'a, A: BlockAllocator> Deref for BlockIndirect<'a, A> {
    type Target = [BlockData<'a, A>; ENTRY_COUNT];

    fn deref(&self) -> &Self::Target {
        unsafe { &*(self.inner.inner.as_ref() as *const _ as *const _) }
    }
}

impl<A: BlockAllocator> DerefMut for BlockIndirect<'_, A> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { &mut *(self.inner.inner.as_mut() as *mut _ as *mut _) }
    }
}

impl<A: BlockAllocator> Drop for BlockIndirect<'_, A> {
    fn drop(&mut self) {
        if self.is_null() {
            return;
        }

        for item in self.iter_mut() {
            item.set_null();
        }
    }
}

// L2 indirect block
impl<A: BlockAllocator> BlockIndirectIndirect<'_, A> {
    /// Dummy entry representing a missing block
    pub const fn null() -> Self {
        Self {
            inner: BlockRaw::null(),
        }
    }

    /// Allocates a new indirection block
    pub fn new() -> Result<Self, Error> {
        let mut inner = BlockRaw::new()?;
        for item in unsafe { inner.as_uninit_indirect_mut() } {
            item.write(BlockData::null());
        }
        Ok(Self { inner })
    }

    /// Returns `true` if the block this structure refers to has not yet been allocated
    #[inline]
    pub fn is_null(&self) -> bool {
        self.inner.is_null()
    }
}

impl<'a, A: BlockAllocator> Deref for BlockIndirectIndirect<'a, A> {
    type Target = [BlockIndirect<'a, A>; ENTRY_COUNT];

    fn deref(&self) -> &Self::Target {
        unsafe { &*(self.inner.inner.as_ref() as *const _ as *const _) }
    }
}

impl<A: BlockAllocator> DerefMut for BlockIndirectIndirect<'_, A> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { &mut *(self.inner.inner.as_mut() as *mut _ as *mut _) }
    }
}

impl<A: BlockAllocator> Drop for BlockIndirectIndirect<'_, A> {
    fn drop(&mut self) {
        if self.is_null() {
            return;
        }

        for item in self.iter_mut() {
            item.set_null();
        }
    }
}

#[cfg(test)]
mod tests {
    use core::sync::atomic::Ordering;
    use std::vec::Vec;

    use crate::block::{BlockData, BlockIndirect};

    #[test]
    fn block_indirect_allocation() {
        test_allocator_with_counter!(A_COUNTER, A);

        const N: usize = 7;
        const M: usize = 13;

        assert_eq!(A_COUNTER.load(Ordering::Acquire), 0);

        {
            let mut indirect = Vec::new();

            // Allocate indirect blocks
            {
                for _ in 0..N {
                    indirect.push(BlockIndirect::<A>::new().unwrap());
                }
            }

            assert_eq!(A_COUNTER.load(Ordering::Acquire), N);

            // Allocate L1 indirection blocks
            {
                for l1_block in indirect.iter_mut() {
                    for i in 0..M {
                        let l0_block = BlockData::new().unwrap();

                        l1_block[i] = l0_block;
                    }
                }
            }

            // N * M data blocks and N indirection blocks
            assert_eq!(A_COUNTER.load(Ordering::Acquire), N * M + N);

            // Drop 1 indirect block for test
            indirect.pop();

            assert_eq!(A_COUNTER.load(Ordering::Acquire), (N - 1) * M + (N - 1));
        }

        assert_eq!(A_COUNTER.load(Ordering::Acquire), 0);
    }

    #[test]
    fn block_allocation() {
        test_allocator_with_counter!(A_COUNTER, A);

        const N: usize = 13;
        {
            assert_eq!(A_COUNTER.load(Ordering::Acquire), 0);

            {
                let mut s = Vec::new();

                for _ in 0..N {
                    let mut block = BlockData::<A>::new().unwrap();
                    block.fill(1);
                    s.push(block);
                }
                assert_eq!(A_COUNTER.load(Ordering::Acquire), N);
            }

            assert_eq!(A_COUNTER.load(Ordering::Acquire), 0);
        }
    }
}
