use core::{
    borrow::{Borrow, BorrowMut},
    cmp::Ordering,
    mem::{self, ManuallyDrop, MaybeUninit},
    ops::{Deref, DerefMut},
};

use alloc::boxed::Box;
use core::fmt;

/// A small box that can store a value inline if the size and alignment of T is
/// less than or equal to the size and alignment of a boxed type. Typically this
/// will be `sizeof::<usize>()` bytes, but might be larger if
/// `sizeof::<Box<T>>()` is larger than that, like it is for dynamically sized
/// types like `[T]` or `dyn Trait`.
#[derive(Debug)]
#[repr(transparent)]
// We use a box here because a box can be unboxed, while a pointer cannot.
pub struct SmallBox<T>(pub MaybeUninit<Box<T>>);

impl<T: fmt::Display> fmt::Display for SmallBox<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        (**self).fmt(f)
    }
}

impl<T: Ord> Ord for SmallBox<T> {
    fn cmp(&self, other: &Self) -> Ordering {
        self.as_ref().cmp(other.as_ref())
    }
}

impl<T: PartialOrd> PartialOrd for SmallBox<T> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.as_ref().partial_cmp(other.as_ref())
    }
}

impl<T: Eq> Eq for SmallBox<T> {}

impl<T: PartialEq> PartialEq for SmallBox<T> {
    fn eq(&self, other: &Self) -> bool {
        self.as_ref().eq(other.as_ref())
    }
}

impl<T: Default> Default for SmallBox<T> {
    fn default() -> Self {
        Self::new(Default::default())
    }
}

impl<T: Clone> Clone for SmallBox<T> {
    fn clone(&self) -> Self {
        Self::new(self.as_ref().clone())
    }
}

impl<T> Deref for SmallBox<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        self.as_ref()
    }
}

impl<T> DerefMut for SmallBox<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.as_mut()
    }
}

impl<T> AsRef<T> for SmallBox<T> {
    fn as_ref(&self) -> &T {
        Self::as_ref(self)
    }
}
impl<T> AsMut<T> for SmallBox<T> {
    fn as_mut(&mut self) -> &mut T {
        Self::as_mut(self)
    }
}

impl<T> Borrow<T> for SmallBox<T> {
    fn borrow(&self) -> &T {
        &**self
    }
}
impl<T> BorrowMut<T> for SmallBox<T> {
    fn borrow_mut(&mut self) -> &mut T {
        &mut **self
    }
}

impl<T> SmallBox<T> {
    /// must only be called once. takes a reference so this can be called in
    /// drop()
    unsafe fn get_unchecked(&self, inline: bool) -> T {
        if inline {
            unsafe { mem::transmute_copy::<MaybeUninit<Box<T>>, T>(&self.0) }
        } else {
            unsafe { *self.0.assume_init_read() }
        }
    }

    pub fn as_ref(&self) -> &T {
        unsafe {
            if Self::is_inline() {
                mem::transmute::<&MaybeUninit<Box<T>>, &T>(&self.0)
            } else {
                self.0.assume_init_ref()
            }
        }
    }

    pub fn as_mut(&mut self) -> &mut T {
        unsafe {
            if Self::is_inline() {
                mem::transmute::<&mut MaybeUninit<Box<T>>, &mut T>(&mut self.0)
            } else {
                self.0.assume_init_mut()
            }
        }
    }

    pub fn into_inner(self) -> T {
        let this = ManuallyDrop::new(self);
        let inline = Self::is_inline();

        // SAFETY: inline is correctly calculated and this function
        // consumes `self`
        unsafe { this.get_unchecked(inline) }
    }

    #[inline(always)]
    pub const fn is_inline() -> bool {
        // the value can be stored inline iff the size of T is equal or
        // smaller than the size of the boxed type and the alignment of the
        // boxed type is an integer multiple of the alignment of T
        (size_of::<Box<MaybeUninit<T>>>() >= size_of::<MaybeUninit<T>>())
            & (align_of::<Box<MaybeUninit<T>>>() >= align_of::<MaybeUninit<T>>())
    }

    pub fn new(value: T) -> Self {
        let inline = Self::is_inline();

        if inline {
            let mut this = MaybeUninit::new(Self(MaybeUninit::uninit()));
            unsafe {
                this.as_mut_ptr().cast::<T>().write(value);
                this.assume_init()
            }
        } else {
            Self(MaybeUninit::new(Box::new(value)))
        }
    }
}

impl<T> Drop for SmallBox<T> {
    fn drop(&mut self) {
        unsafe {
            if Self::is_inline() {
                mem::transmute::<_, &mut MaybeUninit<T>>(&mut self.0).assume_init_drop();
            } else {
                self.0.assume_init_drop();
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use core::sync::atomic::AtomicUsize;

    use super::*;

    #[test]
    fn value_inline() {
        assert!(SmallBox::<u32>::is_inline(), "u32 should be inline");
        assert!(SmallBox::<u8>::is_inline(), "u8 should be inline");
        assert!(
            SmallBox::<Box<u32>>::is_inline(),
            "Box<u32> should be inline"
        );
        assert!(
            SmallBox::<[u32; 2]>::is_inline(),
            "[u32; 2] should be inline"
        );
        assert!(
            !SmallBox::<[u32; 3]>::is_inline(),
            "[u32; 3] should not be inline"
        );
        assert!(SmallBox::<usize>::is_inline(), "usize should be inline");

        #[repr(C, align(16))]
        struct LargeType(u8);
        assert!(
            !SmallBox::<LargeType>::is_inline(),
            "LargeType should not be inline"
        );

        #[repr(C, align(4))]
        struct SmallType(u8);
        assert!(
            SmallBox::<SmallType>::is_inline(),
            "SmallType should be inline"
        );
    }

    #[test]
    fn unbox_smallbox() {
        let small_box = SmallBox::new(42u32);
        assert_eq!(*small_box, 42);
        let value: u32 = small_box.into_inner();
        assert_eq!(value, 42);
    }

    #[test]
    fn drop_smallbox() {
        struct DropCounter<'a> {
            count: &'a AtomicUsize,
        }
        impl<'a> Drop for DropCounter<'a> {
            fn drop(&mut self) {
                self.count
                    .fetch_add(1, core::sync::atomic::Ordering::Relaxed);
            }
        }

        let drop_count = AtomicUsize::new(0);
        {
            let _small_box = SmallBox::new(DropCounter { count: &drop_count });
            assert_eq!(drop_count.load(core::sync::atomic::Ordering::Relaxed), 0);
        }
        assert_eq!(
            drop_count.load(core::sync::atomic::Ordering::Relaxed),
            1,
            "DropCounter should have been dropped"
        );
    }
}