plumers 1.0.2

use std::{
    error::Error,
    ffi::c_void,
    fmt::Display,
    marker::PhantomData,
    mem::MaybeUninit,
    num::NonZeroUsize,
    ops::{Deref, DerefMut},
    ptr::NonNull,
};

use crate::{
    color::ColorFmt,
    image::{AlphaMode, Image, ImageDest, ImageFormat, Metadata, MetadataMut},
};

// Basically a `std::ptr::Unique<plum_image>`, if that was stabilized.
#[derive(Debug)]
pub(super) struct ImgWrapper(
    NonNull<libplum_sys::plum_image>,
    PhantomData<libplum_sys::plum_image>,
);
// SAFETY: no one besides us has the raw pointer, so we can safely transfer the wrapper to another thread.
//         This is only the case because all of the data contained within the image is `Send` itself
//         and the image is fully self-contained.
unsafe impl Send for ImgWrapper {}
// SAFETY: `ImgWrapper` acts as a container for a unique (non-aliased) `plum_image`, and does not
//         contain any interior mutability, so it is safe to race *reads* from it.
unsafe impl Sync for ImgWrapper {}

pub trait PlumWrapper<Fmt: ColorFmt> {
    fn as_img(&self) -> &libplum_sys::plum_image;
    fn as_img_mut(&mut self) -> &mut libplum_sys::plum_image;
    fn palette(&self) -> Option<&[Fmt]>;
    fn pixel_ref(&self, frame_idx: usize, x: usize, y: usize) -> &Fmt;
    fn pixel_mut(&mut self, frame_idx: usize, x: usize, y: usize) -> &mut Fmt;

    /// # Safety
    ///
    /// This function must only be called with the `T` matching the image's "pixel array" type (`Fmt` for `DirectImage`, `u8` for `PalettedImage`.
    unsafe fn pix_array<T>(&self) -> &[T]
    where
        Self: Sized,
    {
        let img = self.as_img();
        let len = nb_pixels(img);
        let ptr = img.data.cast();
        // SAFETY: the slice is initialized, and the pointer is suitable for the `T` guaranteed by the caller.
        unsafe { std::slice::from_raw_parts(ptr, len) }
    }
    /// # Safety
    ///
    /// This function must only be called with the `T` matching the image's "pixel array" type (`Fmt` for `DirectImage`, `u8` for `PalettedImage`.
    unsafe fn pix_array_mut<T>(&mut self) -> &mut [T]
    where
        Self: Sized,
    {
        let img = self.as_img();
        let len = nb_pixels(img);
        let ptr = img.data.cast();
        // SAFETY: the slice is initialized, and the pointer is suitable for the `T` guaranteed by the caller.
        unsafe { std::slice::from_raw_parts_mut(ptr, len) }
    }
}

impl ImgWrapper {
    /// Creates a new, **uninitialized** allocation bound to the wrapped `plum_image`.
    fn alloc<T>(&mut self, nb_elems: usize) -> Option<NonNull<T>> {
        let alloc_size = nb_elems.checked_mul(std::mem::size_of::<T>()).unwrap();
        assert!(alloc_size <= isize::MAX as usize);
        // It's okay if the allocation size is zero, because libplum adds a header and so will never pass a size of zero to the system allocator.

        // SAFETY: the first arg points to an image,
        //         and the allocation size is only bounded by the type.
        let raw_ptr: *mut c_void = unsafe { libplum_sys::plum_malloc(self.0.as_ptr(), alloc_size) };
        // We are casting a raw `void *` (allocator's agnostic "memory" type) into the "real" type.
        // The system allocator guarantees that it always returns a pointer suitably aligned for any type.
        NonNull::new(raw_ptr.cast())
    }

    /// Resizes an allocation bound to the wrapped `plum_image`.
    ///
    /// # Safety
    ///
    /// - `alloc` must have been allocated with this image.
    /// - This may move the contents of `alloc`.
    unsafe fn realloc<T>(&mut self, alloc: *mut T, nb_elems: usize) -> Option<NonNull<T>> {
        let alloc_size = nb_elems.checked_mul(std::mem::size_of::<T>()).unwrap();
        assert!(alloc_size <= isize::MAX as usize);
        // It's okay if the allocation size is zero, because libplum adds a header and so will never pass a size of zero to the system allocator.

        let raw_ptr: *mut c_void =
            // SAFETY: the first arg points to an image,
            //         the second pointer was allocated by a `plum_malloc` on the same image,
            //         and the allocation size is only bounded by the type.
            unsafe { libplum_sys::plum_realloc(self.0.as_ptr(), alloc.cast(), alloc_size) };
        // We are casting a raw `void *` (allocator's agnostic "memory" type) into the "real" type.
        // The system allocator guarantees that it always returns a pointer suitably aligned for any type.
        NonNull::new(raw_ptr.cast())
    }

    /// Creates and **mostly** initializes a [`plum_image`][libplum_sys::plum_image].
    ///
    /// What is left to initialise:
    /// - `frames` (zero)
    /// - `height` (zero)
    /// - `width` (zero)
    /// - `data` (NULL)
    fn make_inner<Fmt: ColorFmt>() -> Self {
        // SAFETY: no particular precautions.
        let raw = unsafe { libplum_sys::plum_new_image() };
        // SAFETY: `plum_new_image` returns an init'd struct on success (zero'd, except the allocator).
        //         Note that zero-init'd means "default", not `memset(0)`; this assumes that C's
        //         NULL is the same as Rust's NULL (so the pointers are valid), which should be OK.
        //         All of the struct's members being `Copy`, mostly all that matters is that they
        //         are init'd; since they are C types, them being zero-init'd should mean the
        //         values are all valid(?)
        //         As for the lifetime, `plum_new_image` returns a `malloc`'d(-ish) pointer.
        let mut raw = Self::from(NonNull::new(raw).unwrap());
        // Perform a few sanity checks for the header members we're not going to initialise.
        // Ideally we'd check that they are all initialised before the above ref cast, to abort
        // before creating any UB, but that's not really doable. (Let's defer that to Miri, I s'pose.)
        debug_assert_eq!(raw.type_, libplum_sys::PLUM_IMAGE_NONE as _);
        debug_assert_eq!(raw.max_palette_index, 0); // The value itself doesn't really matter, since `palette` is null, but if it is different then it's a symptom of some problem...
        debug_assert_eq!(raw.metadata, std::ptr::null_mut()); // No metadata yet.
        debug_assert_eq!(raw.palette, std::ptr::null_mut()); // No palette.
        debug_assert_eq!(raw.userdata, std::ptr::null_mut()); // No userdata.
        debug_assert_ne!(raw.allocator, std::ptr::null_mut()); // Only statically-allocated images should have a null allocator.

        raw.color_format = Fmt::raw_constant() as _; // The actual values all fit in that type.

        raw
    }

    fn make_fixed<Fmt: ColorFmt, F: FnOnce(NonNull<Fmt>, usize)>(
        nb_frames: u32,
        width: u32,
        height: u32,
        init: F,
    ) -> Self {
        let mut new = Self::make_inner::<Fmt>();
        new.frames = nb_frames;
        new.height = height;
        new.width = width;

        let nb_pixels = nb_pixels(&new);
        // SAFETY: the image's allocator is suitably initialised, so we can call `plum_malloc`.
        //         Additionally, the image's contents are being modified through a pointer derived
        //         from a pointer derived from a mutable reference to the image itself.
        let pixels = new.alloc::<Fmt>(nb_pixels).unwrap();
        new.data = pixels.as_ptr().cast();
        init(pixels, nb_pixels);

        new
    }

    pub fn new_zeroed<Fmt: ColorFmt>(nb_frames: u32, width: u32, height: u32) -> Self {
        Self::make_fixed::<Fmt, _>(nb_frames, width, height, |ptr, len| {
            // SAFETY: all of the `Fmt`s can be zero-init'd, since all they contain is one integer.
            //         Also, the length is a <number of elements>, not <number of bytes>!
            unsafe { std::ptr::write_bytes(ptr.as_ptr(), 0, len) };
        })
    }

    pub unsafe fn new<Fmt: ColorFmt, F: FnOnce(&mut [MaybeUninit<Fmt>])>(
        nb_frames: u32,
        width: u32,
        height: u32,
        init: F,
    ) -> Self {
        Self::make_fixed(nb_frames, width, height, |ptr: NonNull<Fmt>, len| {
            // SAFETY: the slice is valid (aligned and correctly-sized), just uninitialized.
            //         Hence why the pointer is cast to `MaybeUninit<Fmt>` instead (which is `#[repr(transparent)]`).
            init(unsafe { std::slice::from_raw_parts_mut(ptr.cast().as_ptr(), len) });
        })
    }

    pub fn collect<T, It: IntoIterator<Item = T>>(&mut self, it: It) -> NonNull<[T]> {
        let elem_size = std::mem::size_of::<T>();
        let mut iter = it.into_iter();

        let mut len = iter.size_hint().0;
        let mut size = len.checked_mul(elem_size).unwrap(); // Avoid allocation size overflow.
        assert!(size <= isize::MAX as usize); // Allocations larger than this are UB in Rust.
        let mut ptr =
            // SAFETY: the first arg points to an image,
            //         and the allocation size is only bounded by the type.
            NonNull::new(unsafe { libplum_sys::plum_malloc(self.0.as_ptr(), size) }).unwrap();

        // First, read as many elements as the iterator promised to return.
        for elem in {
            // SAFETY: the cast is fine, because `MaybeUninit` is `#[repr(transparent)]`.
            let mut ptr = NonNull::slice_from_raw_parts(ptr.cast::<MaybeUninit<T>>(), len);
            // SAFETY: this reference will not live any longer than the loop,
            //         in which the pointer is not used again.
            unsafe { ptr.as_mut() }
        } {
            elem.write(
                iter.next()
                    .expect("Iterator returned fewer elements than promised!"),
            );
        }

        // Now, the iterator might still yield some more elements.
        for extra in iter {
            size += elem_size; // This cannot overflow, because `isize::MAX * 2 <= usize::MAX`.
            assert!(size <= isize::MAX as usize); // Allocations larger than this are UB in Rust.

            // SAFETY: `ptr` hails from `plum_malloc(self.0)`, so we're reallocating using the same allocator.
            ptr = NonNull::new(unsafe {
                libplum_sys::plum_realloc(self.0.as_ptr(), ptr.as_ptr(), size)
            })
            .unwrap();
            // SAFETY: we are offsetting into a single contiguous allocation.
            let new_elem = unsafe { ptr.as_ptr().add(len) }.cast::<T>();
            // SAFETY: the new element hasn't been initialised yet, so we can safely write to it.
            unsafe { new_elem.write(extra) };
            len += 1;
        }

        NonNull::slice_from_raw_parts(ptr.cast(), len)
    }
}

impl Clone for ImgWrapper {
    fn clone(&self) -> Self {
        // TODO: `plum_copy_image` only performs a shallow copy of `.userdata` if it's non-NULL!
        //       This will become a problem if we allow associating anything with the image...
        debug_assert_eq!(self.deref().userdata, std::ptr::null_mut());

        // SAFETY: `self.0` points to a valid image.
        let copy = unsafe { libplum_sys::plum_copy_image(self.0.as_ptr()) };
        // The returned pointer may be NULL iff allocation fails (or the source is invalid).
        Self(NonNull::new(copy).unwrap(), PhantomData)
    }
}

impl Drop for ImgWrapper {
    fn drop(&mut self) {
        // SAFETY: `plum_destroy_image`'s only requirement is that the pointer points to a valid struct (or is NULL).
        unsafe { libplum_sys::plum_destroy_image(self.0.as_ptr()) };
    }
}

impl Deref for ImgWrapper {
    type Target = libplum_sys::plum_image;

    fn deref(&self) -> &Self::Target {
        // SAFETY: by construction, the contained pointer always points to a valid (heap-alloc'd) image.
        unsafe { self.0.as_ref() }
    }
}
impl DerefMut for ImgWrapper {
    fn deref_mut(&mut self) -> &mut Self::Target {
        // SAFETY: by construction, the contained pointer always points to a valid (heap-alloc'd) image.
        unsafe { self.0.as_mut() }
    }
}

impl From<NonNull<libplum_sys::plum_image>> for ImgWrapper {
    fn from(value: NonNull<libplum_sys::plum_image>) -> Self {
        Self(value, PhantomData)
    }
}

// Implementation of the `Image` trait for types that wrap `plum_image`.

macro_rules! wrapper_impls_image {
    ($(impl<Fmt: ColorFmt> PlumWrapper<Fmt> for $implementer:ty { $($impl_block:tt)* })*) => {$(
        impl<Fmt: ColorFmt> PlumWrapper<Fmt> for $implementer { $($impl_block)* }

        impl<Fmt: ColorFmt> Image<Fmt> for $implementer {
            fn nb_frames(&self) -> usize {
                self.as_img().frames.try_into().unwrap()
            }
            fn width(&self) -> usize {
                self.as_img().width.try_into().unwrap()
            }
            fn height(&self) -> usize {
                self.as_img().height.try_into().unwrap()
            }

            fn format(&self) -> ImageFormat {
                self.as_img().type_.try_into().unwrap()
            }
            fn set_format(&mut self, format: ImageFormat) {
                self.as_img_mut().type_ = format.into();
            }

            fn alpha_mode(&self) -> AlphaMode {
                if self.as_img().color_format & (libplum_sys::PLUM_ALPHA_INVERT as u8) != 0{
                    AlphaMode::ZeroIsTransparent
                } else {
                    AlphaMode::ZeroIsOpaque
                }
            }
            fn set_alpha_mode(&mut self, mode: AlphaMode) {
                if mode==AlphaMode::ZeroIsTransparent {
                    self.as_img_mut().color_format |= (libplum_sys::PLUM_ALPHA_INVERT as u8);
                } else {
                    self.as_img_mut().color_format &= !(libplum_sys::PLUM_ALPHA_INVERT as u8);
                }
            }

            fn palette(&self) -> Option<&[Fmt]> {
                PlumWrapper::palette(self)
            }
            fn pixel(&self, frame_idx: usize, x: usize, y: usize) -> Fmt {
                *PlumWrapper::pixel_ref(self, frame_idx, x, y)
            }
            fn pixel_mut(&mut self, frame_idx: usize, x: usize, y: usize) -> &mut Fmt {
                PlumWrapper::pixel_mut(self, frame_idx, x, y)
            }

            fn store<Dest: ImageDest>(&self, output: Dest) -> std::io::Result<NonZeroUsize> {
                output.store::<Fmt>(self.as_img())
            }

            fn metadata(&self) -> Metadata<'_> {
                // SAFETY: the pointer is guaranteed by libplum to point to a valid metadata node,
                //         valid for the same lifetime as the image, which the return type binds it to.
                Metadata::new(self.as_img())
            }
            fn metadata_mut(&mut self) -> MetadataMut<'_> {
                // SAFETY: the pointer is guaranteed by libplum to point to a valid metadata node,
                //         valid for the same lifetime as the image, which the return type binds it to.
                MetadataMut::new(self.as_img_mut())
            }
        }
    )*};
}

wrapper_impls_image! {
    impl<Fmt: ColorFmt> PlumWrapper<Fmt> for super::DirectImage<Fmt> {
        fn as_img(&self) -> &libplum_sys::plum_image {
            &self.0
        }
        fn as_img_mut(&mut self) -> &mut libplum_sys::plum_image {
            &mut self.0
        }

        fn palette(&self) -> Option<&[Fmt]> {
            debug_assert_eq!(self.as_img().palette, std::ptr::null_mut());
            None
        }

        fn pixel_ref(&self, frame_idx: usize, x: usize, y: usize) -> &Fmt {
            let nb_rows = frame_idx * self.nb_frames() + y;
            let offset = nb_rows * self.width() + x;

            // SAFETY: all `ColorFmt`s are `#[repr(transparent)]`, and libplum guarantees that the array contains the underlying type.
            let pix_array = unsafe { self.pix_array() };
            &pix_array[offset]
        }
        fn pixel_mut(&mut self, frame_idx: usize, x: usize, y: usize) -> &mut Fmt {
            let nb_rows = frame_idx * self.nb_frames() + y;
            let offset = nb_rows * self.width() + x;

            // SAFETY: all `ColorFmt`s are `#[repr(transparent)]`, and libplum guarantees that the array contains the underlying type.
            let pix_array = unsafe { self.pix_array_mut() };
            &mut pix_array[offset]
        }
    }

    impl<Fmt: ColorFmt> PlumWrapper<Fmt> for super::PalettedImage<Fmt> {
        fn as_img(&self) -> &libplum_sys::plum_image {
            &self.0
        }
        fn as_img_mut(&mut self) -> &mut libplum_sys::plum_image {
            &mut self.0
        }

        fn palette(&self) -> Option<&[Fmt]> {
            Some(self.palette())
        }

        fn pixel_ref(&self, frame_idx: usize, x: usize, y: usize) -> &Fmt {
            let nb_rows = frame_idx * self.nb_frames() + y;
            let offset = nb_rows * self.width() + x;

            // SAFETY: ibplum guarantees that the array contains `u8`s.
            let idx_array = unsafe { self.pix_array::<u8>() };

            let idx = idx_array[offset];
            &self.palette()[usize::from(idx)]
        }
        fn pixel_mut(&mut self, frame_idx: usize, x: usize, y: usize) -> &mut Fmt {
            let nb_rows = frame_idx * self.nb_frames() + y;
            let offset = nb_rows * self.width() + x;

            // SAFETY: ibplum guarantees that the array contains `u8`s.
            let idx_array = unsafe { self.pix_array::<u8>() };

            let idx = idx_array[offset];
            &mut self.palette_mut()[usize::from(idx)]
        }
    }

    impl<Fmt: ColorFmt> PlumWrapper<Fmt> for super::DynImage<Fmt> {
        fn as_img(&self) -> &libplum_sys::plum_image {
            match self {
                Self::Direct(image) => PlumWrapper::as_img(image),
                Self::Paletted(image) => PlumWrapper::as_img(image),
            }
        }
        fn as_img_mut(&mut self) -> &mut libplum_sys::plum_image {
            match self {
                Self::Direct(image) => PlumWrapper::as_img_mut(image),
                Self::Paletted(image) => PlumWrapper::as_img_mut(image),
            }
        }

        fn palette(&self) -> Option<&[Fmt]> {
            match self {
                Self::Direct(image) => {
                    debug_assert_eq!(image.as_img().palette, std::ptr::null_mut());
                    PlumWrapper::palette(image)
                }
                Self::Paletted(image) => {
                    debug_assert_ne!(image.as_img().palette, std::ptr::null_mut());
                    PlumWrapper::palette(image)
                }
            }
        }

        fn pixel_ref(&self, frame_idx: usize, x: usize, y: usize) -> &Fmt {
            match self {
                Self::Direct(image) => PlumWrapper::pixel_ref(image, frame_idx, x, y),
                Self::Paletted(image) => PlumWrapper::pixel_ref(image, frame_idx, x, y),
            }
        }
        fn pixel_mut(&mut self, frame_idx: usize, x: usize, y: usize) -> &mut Fmt {
            match self {
                Self::Direct(image) => PlumWrapper::pixel_mut(image, frame_idx, x, y),
                Self::Paletted(image) => PlumWrapper::pixel_mut(image, frame_idx, x, y),
            }
        }
    }
}

// Little helper function for the above.
fn nb_pixels(image: &libplum_sys::plum_image) -> usize {
    usize::try_from(image.frames).unwrap()
        * usize::try_from(image.height).unwrap()
        * usize::try_from(image.width).unwrap()
}

macro_rules! convert_to_from {
    ($($t:ty),* $(,)?) => {$(
        impl From<ImageFormat> for $t {
            fn from(value: ImageFormat) -> Self {
                value as $t
            }
        }

        impl TryFrom<$t> for ImageFormat {
            type Error = RawImgFmtOutOfRange;

            fn try_from(value: $t) -> Result<Self, Self::Error> {
                match value.into() {
                    libplum_sys::PLUM_IMAGE_NONE => Ok(Self::None),
                    libplum_sys::PLUM_IMAGE_BMP => Ok(Self::Bmp),
                    libplum_sys::PLUM_IMAGE_GIF => Ok(Self::Gif),
                    libplum_sys::PLUM_IMAGE_PNG => Ok(Self::Png),
                    libplum_sys::PLUM_IMAGE_APNG => Ok(Self::Apng),
                    libplum_sys::PLUM_IMAGE_JPEG => Ok(Self::Jpeg),
                    libplum_sys::PLUM_IMAGE_PNM => Ok(Self::Pnm),
                    _ => Err(RawImgFmtOutOfRange),
                }
            }
        }
    )*};
}
convert_to_from!(libplum_sys::plum_image_types, u8, u16);

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct RawImgFmtOutOfRange;
impl Display for RawImgFmtOutOfRange {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str("raw `plum_image_types` out of range")
    }
}
impl Error for RawImgFmtOutOfRange {}