use std::mem::MaybeUninit;

#[repr(align(64))]
struct Inline<const CAP: usize> {
    _bytes: [u8; CAP],
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum Mode {
    Empty,
    Inline,
    Heap,
}

pub(super) struct SmallSlot<const CAP: usize> {
    inline: MaybeUninit<Inline<CAP>>,
    ptr: *mut u8,
    drop: unsafe fn(*mut u8),
    mode: Mode,
}

impl<const CAP: usize> Default for SmallSlot<CAP> {
    fn default() -> Self {
        Self::new()
    }
}

impl<const CAP: usize> SmallSlot<CAP> {
    #[must_use]
    pub(super) fn new() -> Self {
        Self {
            inline: MaybeUninit::uninit(),
            ptr: std::ptr::null_mut(),
            drop: Self::drop_nop,
            mode: Mode::Empty,
        }
    }

    pub(super) fn clear(&mut self) {
        if self.mode == Mode::Empty {
            return;
        }

        unsafe { (self.drop)(self.ptr) };
        self.ptr = std::ptr::null_mut();
        self.drop = Self::drop_nop;
        self.mode = Mode::Empty;
    }

    pub(super) fn put<T: 'static>(&mut self, value: T) {
        self.clear();
        if std::mem::size_of::<T>() <= CAP
            && std::mem::align_of::<T>() <= std::mem::align_of::<Inline<CAP>>()
        {
            let ptr = self.inline.as_mut_ptr().cast::<u8>();

            unsafe { (ptr as *mut T).write(value) };
            self.ptr = ptr;
            self.drop = Self::drop_inline::<T>;
            self.mode = Mode::Inline;
            return;
        }
        let raw = Box::into_raw(Box::new(value)) as *mut u8;
        self.ptr = raw;
        self.drop = Self::drop_heap::<T>;
        self.mode = Mode::Heap;
    }

    unsafe fn drop_inline<T>(raw: *mut u8) {
        unsafe { std::ptr::drop_in_place(raw as *mut T) };
    }

    unsafe fn drop_heap<T>(raw: *mut u8) {
        unsafe {
            drop(Box::from_raw(raw as *mut T));
        }
    }

    unsafe fn drop_nop(_raw: *mut u8) {}
}

impl<const CAP: usize> Drop for SmallSlot<CAP> {
    fn drop(&mut self) {
        self.clear();
    }
}
