plumers 1.0.2

use std::{
    ffi::{c_int, c_void},
    fs::File,
    io::{BufReader, Read},
    mem::MaybeUninit,
    path::Path,
    ptr::NonNull,
};

use crate::{color::ColorFmt, Error};

use super::AlphaMode;

/// Parameters affecting how an image should be loaded.
#[derive(Debug, Clone, Copy, Default)]
pub struct LoadFlags {
    /// If `true`, the image's alpha channel is treated as fully opaque.
    ///
    /// (And thus stored as 0x00 or 0xFF depending on [`AlphaMode`].)
    pub remove_alpha: bool,
    /// How palettes should be sorted.
    /// Ignored if palettes are not used (e.g. loading a [`DirectImage`][super::DirectImage]).
    ///
    /// Note that unless [`sort_existing`][Self::sort_existing] is enabled, only generated palettes get sorted!
    pub palette_sort: PaletteSort,
    /// If `false` (the default), and the loaded image contains an embedded colour palette, it will be loaded as-is.
    /// Ignored if palettes are not used (e.g. loading a [`DirectImage`][super::DirectImage]).
    pub sort_existing: bool,
    /// If `true`, and the loaded image contains an embedded colour palette, duplicate colours will be merged together, and unused colours removed.
    /// Ignored if palettes are not used (e.g. loading a [`DirectImage`][super::DirectImage]).
    ///
    /// Note that generated colour palettes are always "reduced".
    pub reduce_palette: bool,
}

/// Select how generated palettes should be sorted.
///
/// Sorting is performed on the following key, after upscaling all channels to 16-bit:
///
/// ```rs
/// (red * 299 + green * 587 + blue * 114, red + green + blue, alpha >> 7)
/// ```
///
/// The first component approximates luminance (and is **the only one inverted** by `DarkFirst`);
/// the second breaks ties in favour of colours with higher components;
/// and the last one ensures that ties are broken by moving more transparent colours last.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
pub enum PaletteSort {
    /// Prefer brighter colours first. (This is the default.)
    #[default]
    LightFirst,
    /// Prefer darker colours first. (But the tie-breaker is unchanged!)
    DarkFirst,
}

/// An object that an image can be loaded from.
///
/// This trait is designed to be implementable outside of this crate, and could in theory be used outside of it.
/// But since [`load()`][Self::load()] returns a raw [`plum_image`][libplum_sys::plum_image], it is not very useful outside of this crate.
/// It is mainly provided for convenience, for any code that wishes to be generic over accepted sources.
///
/// # Provided impls
///
/// The crate provides three classes of `impl`s:
/// - [`Path`] (and [`str`], which is path-like): opens the named file, and loads from it.
/// - Anything that implements [`Read`] (see below for details).
/// - Callbacks, which are the most flexible but also the most verbose.
///
/// ## [`Read`]ables
///
/// It is possible for a type to implement both [`FnMut(&mut [u8]) -> io::Result<u16>`][FnMut] (which qualifies it as a callback) and [`Read`] (which qualifies it as a readable).
///
/// For the purpose of disambiguation (which has to be done all the time because of Rust's [orphan rules]), "readables" have to be wrapped in a newtype: [`Input`].
/// Please remember that if `T` implements [`Read`], so does `&mut T` (and [`Read::by_ref()`] is a convenience to help with that).
///
/// However, some readables have "shortcut" impls provided for convenience (so you can use [`File`] directly instead of `Input<File>`, for example).
/// They are strictly developer-convenience shortcuts, with one exception: [`&[u8]`][slice] (and, by extension, [`&mut [u8]`][slice]), which perform *better* if not wrapped in `Input`, for internal reasons.
///
/// In fact, "readables" are implemented in terms of callbacks, and provided for convenience's sake, since they're so common (and writing the callback each time would get tedious).
///
/// ## Callbacks
///
/// Trying to use a closure as callback may yield error `E0308`, `E0599`, or other errors about trait bounds:
///
/// ```rust,compile_fail
/// # use std::io::Read;
/// use plumers::prelude::*;
///
/// # fn blah(file: &mut std::fs::File, flags: LoadFlags, alpha_invert: AlphaMode) -> std::io::Result<DirectImage32> {
/// let callback = |buffer| file.read(buffer).map(|len| len as u16);
/// DirectImage32::load(callback, flags, alpha_invert)
/// # }
/// ```
///
/// ```text
/// error[E0308]: mismatched types
///  --> plumers/src/load.rs:9:1
///   |
/// 9 | DirectImage32::load(callback, flags, alpha_invert)
///   | ^^^^^^^^^^^^^^^^^^^ one type is more general than the other
///   |
///   = note: expected trait `for<'a> FnMut(&'a mut [u8])`
///              found trait `FnMut(&mut [u8])`
/// note: this closure does not fulfill the lifetime requirements
///  --> plumers/src/load.rs:8:16
///   |
/// 8 | let callback = |buffer| file.read(buffer).map(|len| len as u16);
///   |                ^^^^^^^^
/// help: consider specifying the type of the closure parameters
///   |
/// 8 | let callback = |buffer: &_| file.read(buffer).map(|len| len as u16);
///   |                ~~~~~~~~~~~~
/// ```
///
/// The fix is to explicitly specify the parameter's type:
///
/// ```rust
/// # use std::io::Read;
/// use plumers::prelude::*;
///
/// # fn blah(file: &mut std::fs::File, flags: LoadFlags, alpha_invert: AlphaMode) -> std::io::Result<DirectImage32> {
/// //                    ↓↓↓↓↓↓↓↓ add this
/// let callback = |buffer: &mut _| file.read(buffer).map(|len| len as u16);
/// DirectImage32::load(callback, flags, alpha_invert)
/// # }
/// ```
///
/// (The implicit type's lifetime seems to be less general than the explicit type's lifetime, despite both being elided.)
///
/// [orphan rules]: <https://doc.rust-lang.org/stable/reference/items/implementations.html#orphan-rules>
pub trait ImageSource {
    /// Attempts to load a [`plum_image`][libplum_sys::plum_image] from this source.
    fn load<Fmt: ColorFmt>(
        self,
        flags: LoadFlags,
        alpha_invert: AlphaMode,
        palette_mode: PaletteMode,
        mem_limit: usize,
    ) -> std::io::Result<NonNull<libplum_sys::plum_image>>;
}

/// Please read [the trait documentation](#callbacks) if you are having trouble using this impl.
impl<F: FnMut(&mut [u8]) -> std::io::Result<u16>> ImageSource for F {
    fn load<Fmt: ColorFmt>(
        self,
        flags: LoadFlags,
        alpha_invert: AlphaMode,
        palette_mode: PaletteMode,
        mem_limit: usize,
    ) -> std::io::Result<NonNull<libplum_sys::plum_image>> {
        struct Userdata<E, F>(F, MaybeUninit<E>);

        let mut userdata = Userdata(self, MaybeUninit::<std::io::Error>::uninit());
        let callback_struct = libplum_sys::plum_callback {
            callback: Some(callback_wrapper::<F>),
            userdata: &mut userdata as *mut _ as *mut c_void,
        };
        extern "C" fn callback_wrapper<F: FnMut(&mut [u8]) -> std::io::Result<u16>>(
            userdata: *mut c_void,
            buffer: *mut c_void,
            buf_len: c_int,
        ) -> c_int {
            // First, round-trip the `callback` pointer back to the same as the outer `callback` argument.
            let userdata = userdata.cast::<Userdata<_, F>>();
            // SAFETY: the pointer is extracted from the outer function's `&mut userdata`, so it's safe to deref.
            let Userdata(callback, error_return) = unsafe { &mut *userdata };

            // Then, prepare the buffer.

            // SAFETY: we are just undoing C's type-erasure.
            let buffer = buffer.cast();
            // "The `buf_len` argument is the size of the supplied buffer, and it will always be positive and no larger than `0x7fff`."
            let buf_len = buf_len as u16;
            // Rust semantics require the buffer to be initialised before creating the slice.
            // TODO: is there a better way to do the init?
            // SAFETY: the buffer length is the associated length returned by libplum.
            unsafe { std::ptr::write_bytes(buffer, 0, buf_len.into()) };
            // SAFETY: the buffer length is the associated length returned by libplum, and we just init'd the buffer.
            let buffer = unsafe { std::slice::from_raw_parts_mut(buffer, buf_len.into()) };

            // Then, simply call it!
            let ret = callback(buffer);

            match ret {
                Ok(size) => {
                    assert!(size <= buf_len, "Read more bytes than were available?!");
                    size.into()
                }
                Err(err) => {
                    // SAFETY: since we return a negative value, libplum will not call us again;
                    //         therefore, we can't be overwriting something already initialised.
                    error_return.write(err);
                    -1
                }
            }
        }

        let ptr = &callback_struct as *const _;
        let raw_flags = flags.to_raw::<Fmt>(alpha_invert, palette_mode);
        let mut error: std::ffi::c_uint = 0;
        // Rust does not handle allocations larger than that.
        assert!(mem_limit <= max_nb_pixels::<Fmt>());
        // SAFETY:
        //  - `ptr` points to a C-ABI function
        //  - The `mode` argument is CALLBACK
        //  - The memory allocation will not overflow `isize::max()` bytes thanks to the above `assert!`
        let image = unsafe {
            #[allow(trivial_numeric_casts)] // Necessary on some platforms, superfluous on others.
            libplum_sys::plum_load_image_limited(
                ptr as *const c_void,
                libplum_sys::PLUM_MODE_CALLBACK,
                raw_flags as _,
                mem_limit,
                &mut error as _,
            )
        };
        // SAFETY: `plum_load_image`'s API is to return a pointer to a struct on success, and NULL on error.
        //         - The pointer comes from `malloc`, which ensures alignment.
        //         - The correct size was passed to `malloc`, which ensures it's dereferenceable.
        //         - `plum_new_image` does a struct-to-struct assignment, which inits the entire struct
        //           (except padding, but AFAIK that's fine?).
        //         - The original pointer is lost, and no copy exists(*), which guarantees no aliasing.
        // (*) Actually, libplum stores a copy of the pointer in its "allocator", but does not
        //     make use of it outside of `plum_destroy_image`.
        NonNull::new(image).ok_or_else(|| {
            #[allow(trivial_numeric_casts)] // Necessary on some platforms, superfluous on others.
            let error = Error::from_raw(error as _);
            std::io::Error::new(std::io::ErrorKind::Other, error)
        })
    }
}

/// This struct is required to work around Rust's "[orphan rules]".
/// See [`ImageSource`].
///
/// Wrap any [`Read`]able you want to load an image from in this structure.
///
/// [orphan rules]: <https://doc.rust-lang.org/stable/reference/items/implementations.html#orphan-rules>
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Input<R: Read>(pub R);
impl<R: Read> From<R> for Input<R> {
    fn from(value: R) -> Self {
        Self(value)
    }
}

impl<R: Read> ImageSource for Input<R> {
    fn load<Fmt: ColorFmt>(
        mut self,
        flags: LoadFlags,
        alpha_invert: AlphaMode,
        palette_mode: PaletteMode,
        mem_limit: usize,
    ) -> std::io::Result<NonNull<libplum_sys::plum_image>> {
        let callback = |buffer: &mut [u8]| {
            self.0
                .read(buffer)
                // The `as u16` truncation is always OK, because libplum gives us a suitably-sized buffer.
                .map(|len| len as u16)
        };
        callback.load::<Fmt>(flags, alpha_invert, palette_mode, mem_limit)
    }
}

macro_rules! read_shortcut {
    ($(
        impl$(< $($ty_param:ident $(: $bound:tt $(+ $bounds:tt)* )? ),+ $(,)? >)?
        ImageSource for $t:ty {}
    )*) => {$(
        /// Convenience shorthand for the generic [`std::io::Read`] impl.
        impl$(< $($ty_param $(: $bound $(+ $bounds)* )? ),+ >)?
        ImageSource for $t {
            fn load<Fmt: ColorFmt>(
                self,
                flags: LoadFlags,
                alpha_invert: AlphaMode,
                palette_mode: PaletteMode,
                mem_limit: usize,
            ) -> std::io::Result<NonNull<libplum_sys::plum_image>> {
                Input(self).load::<Fmt>(flags, alpha_invert, palette_mode, mem_limit)
            }
        }
    )*};
}
read_shortcut! {
    // Note that we cannot use libplum's "file" loading mode, because we cannot extract the file's path, much less convert it to UTF-8.
    impl ImageSource for File {}
    impl ImageSource for &File {}
    impl ImageSource for &mut File {}

    impl<R: Read> ImageSource for BufReader<R> {}
    impl<R: Read> ImageSource for &mut BufReader<R> {}
}

/// This implementation should be preferred over using the generic [`Read`] (or any of its shortcuts), as it has **less** overhead.
impl ImageSource for &[u8] {
    fn load<Fmt: ColorFmt>(
        self,
        flags: LoadFlags,
        alpha_invert: AlphaMode,
        palette_mode: PaletteMode,
        mem_limit: usize,
    ) -> std::io::Result<NonNull<libplum_sys::plum_image>> {
        let buffer = self as *const _ as *const c_void;
        let size = self.len();
        let raw_flags = flags.to_raw::<Fmt>(alpha_invert, palette_mode);

        let mut error: std::ffi::c_uint = 0;
        // Rust does not handle allocations larger than that.
        assert!(mem_limit <= max_nb_pixels::<Fmt>());
        // SAFETY:
        //  - `ptr` points to a C-ABI function
        //  - The `mode` argument is CALLBACK
        //  - The memory allocation will not overflow `isize::max()` bytes thanks to the above `assert!`
        let image = unsafe {
            #[allow(trivial_numeric_casts)] // Necessary on some platforms, superfluous on others.
            libplum_sys::plum_load_image_limited(
                buffer,
                size,
                raw_flags as _,
                mem_limit,
                &mut error as _,
            )
        };
        // SAFETY: `plum_load_image`'s API is to return a pointer to a struct on success, and NULL on error.
        //         - The pointer comes from `malloc`, which ensures alignment.
        //         - The correct size was passed to `malloc`, which ensures it's dereferenceable.
        //         - `plum_new_image` does a struct-to-struct assignment, which inits the entire struct
        //           (except padding, but AFAIK that's fine?).
        //         - The original pointer is lost, and no copy exists(*), which guarantees no aliasing.
        // (*) Actually, libplum stores a copy of the pointer in its "allocator", but does not
        //     make use of it outside of `plum_destroy_image`.
        NonNull::new(image).ok_or_else(|| {
            #[allow(trivial_numeric_casts)] // Necessary on some platforms, superfluous on others.
            let error = Error::from_raw(error as _);
            std::io::Error::new(std::io::ErrorKind::Other, error)
        })
    }
}

/// This implementation should be preferred over using the generic [`Read`] (or any of its shortcuts), as it has **less** overhead.
impl ImageSource for &mut [u8] {
    fn load<Fmt: ColorFmt>(
        self,
        flags: LoadFlags,
        alpha_invert: AlphaMode,
        palette_mode: PaletteMode,
        mem_limit: usize,
    ) -> std::io::Result<NonNull<libplum_sys::plum_image>> {
        ImageSource::load::<Fmt>(self as &[u8], flags, alpha_invert, palette_mode, mem_limit)
    }
}

/// This implementation should be preferred over using the generic [`Read`] (or any of its shortcuts), as it has **less** overhead.
impl<const N: usize> ImageSource for [u8; N] {
    fn load<Fmt: ColorFmt>(
        self,
        flags: LoadFlags,
        alpha_invert: AlphaMode,
        palette_mode: PaletteMode,
        mem_limit: usize,
    ) -> std::io::Result<NonNull<libplum_sys::plum_image>> {
        ImageSource::load::<Fmt>(
            self.as_slice(),
            flags,
            alpha_invert,
            palette_mode,
            mem_limit,
        )
    }
}

/// This implementation should be preferred over using the generic [`Read`] (or any of its shortcuts), as it has **less** overhead.
impl<const N: usize> ImageSource for &[u8; N] {
    fn load<Fmt: ColorFmt>(
        self,
        flags: LoadFlags,
        alpha_invert: AlphaMode,
        palette_mode: PaletteMode,
        mem_limit: usize,
    ) -> std::io::Result<NonNull<libplum_sys::plum_image>> {
        ImageSource::load::<Fmt>(
            self.as_slice(),
            flags,
            alpha_invert,
            palette_mode,
            mem_limit,
        )
    }
}

/// This implementation should be preferred over using the generic [`Read`] (or any of its shortcuts), as it has **less** overhead.
impl<const N: usize> ImageSource for &mut [u8; N] {
    fn load<Fmt: ColorFmt>(
        self,
        flags: LoadFlags,
        alpha_invert: AlphaMode,
        palette_mode: PaletteMode,
        mem_limit: usize,
    ) -> std::io::Result<NonNull<libplum_sys::plum_image>> {
        ImageSource::load::<Fmt>(
            self.as_slice(),
            flags,
            alpha_invert,
            palette_mode,
            mem_limit,
        )
    }
}

impl ImageSource for &Path {
    fn load<Fmt: ColorFmt>(
        self,
        flags: LoadFlags,
        alpha_invert: AlphaMode,
        palette_mode: PaletteMode,
        mem_limit: usize,
    ) -> std::io::Result<NonNull<libplum_sys::plum_image>> {
        let file = File::open(self)?;
        Input(file).load::<Fmt>(flags, alpha_invert, palette_mode, mem_limit)
    }
}

impl ImageSource for &str {
    fn load<Fmt: ColorFmt>(
        self,
        flags: LoadFlags,
        alpha_invert: AlphaMode,
        palette_mode: PaletteMode,
        mem_limit: usize,
    ) -> std::io::Result<NonNull<libplum_sys::plum_image>> {
        let file = File::open(self)?;
        Input(file).load::<Fmt>(flags, alpha_invert, palette_mode, mem_limit)
    }
}

impl LoadFlags {
    /// Encodes `self`, `alpha_invert`, and `palette_mode` into a bitfield suitable for libplum.
    pub fn to_raw<Fmt: ColorFmt>(
        self,
        alpha_invert: AlphaMode,
        palette_mode: PaletteMode,
    ) -> libplum_sys::plum_flags {
        let mut raw = Fmt::raw_constant();
        raw |= alpha_invert.to_raw();
        if self.remove_alpha {
            raw |= libplum_sys::PLUM_ALPHA_REMOVE;
        }
        raw |= palette_mode.to_raw();
        raw |= self.palette_sort.to_raw();
        if self.sort_existing {
            raw |= libplum_sys::PLUM_SORT_EXISTING;
        }
        if self.reduce_palette {
            raw |= libplum_sys::PLUM_PALETTE_REDUCE;
        }
        raw
    }
}

/// Returns how many pixels of a given format an image can contain at most.
///
/// Rust imposes limits on the size of objects, which is where the limit this function indicates stems from.
pub const fn max_nb_pixels<Fmt: ColorFmt>() -> usize {
    isize::MAX as usize / std::mem::size_of::<Fmt::Raw>()
}

// This is largely an implementation detail internal to `ImageSource::load`. No need to list it in the crate docs.
#[doc(hidden)]
/// How palettes should be processed.
#[derive(Debug, Clone, Copy)]
pub enum PaletteMode {
    /// Do not produce a palette, ever.
    /// Turn a paletted image into a "direct" one.
    None,
    /// Load the palette from the image if there is one, and keep the image "direct" otherwise.
    Load,
    /// Load the palette from the image if there is one.
    /// Otherwise, attempt to generate a palette, but don't fret if there are more than 256 unique colours.
    Generate,
    /// Load the palette from the image if there is one, and generate one otherwise.
    /// If the image contains more than 256 unique colours, produce an error.
    Force,
}

impl PaletteMode {
    /// Encodes `self` into a bitfield suitable for libplum.
    pub fn to_raw(self) -> libplum_sys::plum_flags {
        match self {
            PaletteMode::None => libplum_sys::PLUM_PALETTE_NONE,
            PaletteMode::Load => libplum_sys::PLUM_PALETTE_LOAD,
            PaletteMode::Generate => libplum_sys::PLUM_PALETTE_GENERATE,
            PaletteMode::Force => libplum_sys::PLUM_PALETTE_FORCE,
        }
    }
}

impl PaletteSort {
    /// Encodes `self` into a bitfield suitable for libplum.
    pub fn to_raw(self) -> libplum_sys::plum_flags {
        match self {
            PaletteSort::LightFirst => libplum_sys::PLUM_SORT_LIGHT_FIRST,
            PaletteSort::DarkFirst => libplum_sys::PLUM_SORT_DARK_FIRST,
        }
    }
}

impl AlphaMode {
    /// Encodes `self` into a bitfield suitable for libplum.
    pub fn to_raw(self) -> libplum_sys::plum_flags {
        match self {
            AlphaMode::ZeroIsOpaque => 0,
            AlphaMode::ZeroIsTransparent => libplum_sys::PLUM_ALPHA_INVERT,
        }
    }
}