use core::{
    alloc::{AllocError, Allocator, GlobalAlloc, Layout},
    ptr::NonNull,
};

use embassy_sync::blocking_mutex::Mutex;
use rlsf::Tlsf;
use static_cell::ConstStaticCell;

use crate::locks::AllocatorLock;

type Heap = Tlsf<'static, usize, usize, { usize::BITS as usize }, { usize::BITS as usize }>;

struct HeapInner {
    heap: Heap,
}

pub struct PicoHeap<const MEMORY_SIZE: usize> {
    inner: Mutex<AllocatorLock, HeapInner>,
    block: ConstStaticCell<[u8; MEMORY_SIZE]>,
}

#[derive(Debug)]
#[cfg_attr(feature = "defmt", derive(defmt::Format))]
pub enum HeapError {
    MemoryBlockTooSmall,
    BlockAlreadyTaken,
}

impl<const MEMORY_SIZE: usize> PicoHeap<MEMORY_SIZE> {
    /// Creates an empty, uninitialised [`PicoHeap`]. This must be created with a sufficiently sized
    /// `MEMORY_SIZE`, else [`PicoHeap::init`] will return a [`HeapError::MemoryBlockTooSmall`] error
    /// when it is called.
    pub const fn empty() -> Self {
        Self {
            inner: Mutex::new(HeapInner { heap: Heap::new() }),
            block: ConstStaticCell::new([0u8; MEMORY_SIZE]),
        }
    }

    /// Initialises the heap. Should be called only once, else this method will return an error.
    pub fn init(&'static self) -> Result<(), HeapError> {
        // SAFETY: We have checked whether the heap was initialised already, and that the block
        // of memory satisfies our lifetime conditions thanks to the &'static self type of the
        // init() method. We also never call lock_mut reetrantly, so this is safe.
        unsafe {
            self.inner.lock_mut(|inner| {
                let block = self.block.try_take().ok_or(HeapError::BlockAlreadyTaken)?;

                inner
                    .heap
                    .insert_free_block_ptr(NonNull::from_mut(block))
                    .ok_or(HeapError::MemoryBlockTooSmall)?;

                Ok(())
            })
        }
    }

    pub fn alloc_block(&self, layout: Layout) -> Option<NonNull<u8>> {
        // SAFETY: lock_mut is never called reentrantly, therefore this is safe.
        unsafe { self.inner.lock_mut(|inner| inner.heap.allocate(layout)) }
    }

    /// # Safety
    ///
    /// The caller must ensure:
    ///
    /// * `ptr` is a block of memory currently allocated via this allocator and,
    /// * `layout` is the same layout that was used to allocate that block of memory.
    ///
    /// Otherwise the behavior is undefined.
    pub unsafe fn dealloc_block(&self, ptr: NonNull<u8>, layout: Layout) {
        // SAFETY: lock_mut is never called reentrantly, therefore this is safe. Deallocation
        // invariants are upheld by the caller.
        unsafe {
            self.inner
                .lock_mut(|inner| inner.heap.deallocate(ptr, layout.align()));
        }
    }
}

unsafe impl<const MEMORY_SIZE: usize> GlobalAlloc for PicoHeap<MEMORY_SIZE> {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        self.alloc_block(layout)
            .map_or(core::ptr::null_mut(), NonNull::as_ptr)
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        unsafe {
            self.dealloc_block(NonNull::new_unchecked(ptr), layout);
        }
    }
}

unsafe impl<const MEMORY_SIZE: usize> Allocator for PicoHeap<MEMORY_SIZE> {
    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, core::alloc::AllocError> {
        match layout.size() {
            0 => Ok(NonNull::slice_from_raw_parts(NonNull::dangling(), 0)),
            size => self
                .alloc_block(layout)
                .map_or(Err(AllocError), |allocated| {
                    Ok(NonNull::slice_from_raw_parts(allocated, size))
                }),
        }
    }

    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
        if layout.size() != 0 {
            unsafe {
                self.dealloc_block(ptr, layout);
            }
        }
    }
}