path: root/engine/src/image.rs

use std::any::type_name;
use std::fs::File;
use std::io::BufReader;
use std::path::Path;

use image_rs::GenericImageView as _;
use zerocopy::{FromBytes, Immutable};

use crate::color::Color;
use crate::data_types::dimens::Dimens;

#[derive(Debug, Clone)]
pub struct Image
{
    inner: image_rs::DynamicImage,
}

impl Image
{
    pub fn open(path: impl AsRef<Path>) -> Result<Self, Error>
    {
        let buffered_reader =
            BufReader::new(File::open(&path).map_err(Error::ReadFailed)?);

        let image_reader = image_rs::ImageReader::with_format(
            buffered_reader,
            image_rs::ImageFormat::from_path(path)
                .map_err(|_| Error::UnsupportedFormat)?,
        );

        Ok(Self {
            inner: image_reader
                .decode()
                .map_err(|err| Error::DecodeFailed(DecodeError(err)))?,
        })
    }

    pub fn try_from_bytes(
        bytes: &[u8],
        dimens: Dimens<u32>,
        color_type: ColorType,
    ) -> Result<Self, FromBytesError>
    {
        use image_rs::{Luma, LumaA, Rgb, Rgba};

        let inner = match color_type {
            ColorType::L8 => image_buf_from_bytes::<Luma<u8>>(dimens, bytes)?.into(),
            ColorType::La8 => image_buf_from_bytes::<LumaA<u8>>(dimens, bytes)?.into(),
            ColorType::Rgb8 => image_buf_from_bytes::<Rgb<u8>>(dimens, bytes)?.into(),
            ColorType::Rgba8 => image_buf_from_bytes::<Rgba<u8>>(dimens, bytes)?.into(),
            ColorType::L16 => image_buf_from_bytes::<Luma<u16>>(dimens, bytes)?.into(),
            ColorType::La16 => image_buf_from_bytes::<LumaA<u16>>(dimens, bytes)?.into(),
            ColorType::Rgb16 => image_buf_from_bytes::<Rgb<u16>>(dimens, bytes)?.into(),
            ColorType::Rgba16 => image_buf_from_bytes::<Rgba<u16>>(dimens, bytes)?.into(),
            ColorType::Rgb32F => image_buf_from_bytes::<Rgb<f32>>(dimens, bytes)?.into(),
            ColorType::Rgba32F => {
                image_buf_from_bytes::<Rgba<f32>>(dimens, bytes)?.into()
            }
        };

        Ok(Self { inner })
    }

    pub fn from_color(dimens: impl Into<Dimens<u32>>, color: impl Into<Color<u8>>)
    -> Self
    {
        let dimens: Dimens<u32> = dimens.into();

        let color: Color<u8> = color.into();

        Self {
            inner: image_rs::RgbImage::from_pixel(
                dimens.width,
                dimens.height,
                image_rs::Rgb([color.red, color.green, color.blue]),
            )
            .into(),
        }
    }

    pub fn from_color_and_alpha(
        dimens: impl Into<Dimens<u32>>,
        color: impl Into<Color<u8>>,
        alpha: u8,
    ) -> Self
    {
        let dimens: Dimens<u32> = dimens.into();

        let color: Color<u8> = color.into();

        Self {
            inner: image_rs::RgbaImage::from_pixel(
                dimens.width,
                dimens.height,
                image_rs::Rgba([color.red, color.green, color.blue, alpha]),
            )
            .into(),
        }
    }

    pub fn dimensions(&self) -> Dimens<u32>
    {
        self.inner.dimensions().into()
    }

    pub fn color_type(&self) -> ColorType
    {
        self.inner.color().into()
    }

    pub fn color_space_is_srgb(&self) -> bool
    {
        match self.inner.color_space().primaries {
            image_rs::metadata::CicpColorPrimaries::SRgb => true,
            _ => false,
        }
    }

    pub fn to_rgba8(&self) -> Self
    {
        Self { inner: self.inner.to_rgba8().into() }
    }

    pub fn as_bytes(&self) -> &[u8]
    {
        self.inner.as_bytes()
    }
}

/// An enumeration over supported color types and bit depths
#[derive(Copy, PartialEq, Eq, Debug, Clone, Hash)]
#[non_exhaustive]
pub enum ColorType
{
    /// Pixel is 8-bit luminance
    L8,

    /// Pixel is 8-bit luminance with an alpha channel
    La8,

    /// Pixel contains 8-bit R, G and B channels
    Rgb8,

    /// Pixel is 8-bit RGB with an alpha channel
    Rgba8,

    /// Pixel is 16-bit luminance
    L16,

    /// Pixel is 16-bit luminance with an alpha channel
    La16,

    /// Pixel is 16-bit RGB
    Rgb16,

    /// Pixel is 16-bit RGBA
    Rgba16,

    /// Pixel is 32-bit float RGB
    Rgb32F,

    /// Pixel is 32-bit float RGBA
    Rgba32F,
}

impl From<image_rs::ColorType> for ColorType
{
    fn from(color_type: image_rs::ColorType) -> Self
    {
        match color_type {
            image_rs::ColorType::L8 => Self::L8,
            image_rs::ColorType::La8 => Self::La8,
            image_rs::ColorType::Rgb8 => Self::Rgb8,
            image_rs::ColorType::Rgba8 => Self::Rgba8,
            image_rs::ColorType::L16 => Self::L16,
            image_rs::ColorType::La16 => Self::La16,
            image_rs::ColorType::Rgb16 => Self::Rgb16,
            image_rs::ColorType::Rgba16 => Self::Rgba16,
            image_rs::ColorType::Rgb32F => Self::Rgb32F,
            image_rs::ColorType::Rgba32F => Self::Rgba32F,
            _ => {
                panic!("Unrecognized image_rs::ColorType variant");
            }
        }
    }
}

#[derive(Debug, thiserror::Error)]
pub enum Error
{
    #[error("Failed to read file")]
    ReadFailed(#[source] std::io::Error),

    #[error("Failed to decode image")]
    DecodeFailed(DecodeError),

    #[error("Unsupported image format")]
    UnsupportedFormat,
}

#[derive(Debug, thiserror::Error)]
#[error(transparent)]
pub struct DecodeError(image_rs::ImageError);

#[derive(Debug, thiserror::Error)]
pub enum FromBytesError
{
    #[error(
        "The number of bytes provided is not enough to fit a image with the given size"
    )]
    NotEnoughBytesForDimensions,

    #[error("Failed to cast bytes to subpixels")]
    CastToSubPixelsFailed(#[from] CastToSubPixelsError),
}

#[derive(Debug, thiserror::Error)]
pub enum CastToSubPixelsError
{
    #[error(
        "Source address {:?} ({}) isn't a multiple of the alignment of the type {} ({})",
        src,
        src_type_name,
        dst_type_name,
        dst_align
    )]
    Alignment
    {
        src_type_name: &'static str,
        dst_type_name: &'static str,
        src: *const u8,
        dst_align: usize,
    },

    #[error(
        "Source size {} ({}) is incorrect size for type {}",
        src_size,
        src_type_name,
        dst_type_name
    )]
    Size
    {
        src_type_name: &'static str,
        dst_type_name: &'static str,
        src_size: usize,
    },
}

fn image_buf_from_bytes<Pixel>(
    dimens: Dimens<u32>,
    bytes: &[u8],
) -> Result<image_rs::ImageBuffer<Pixel, Vec<Pixel::Subpixel>>, FromBytesError>
where
    Pixel: image_rs::Pixel<Subpixel: FromBytes + Immutable>,
{
    let buf = <[Pixel::Subpixel]>::ref_from_bytes(bytes).map_err(|err| match err {
        zerocopy::CastError::Alignment(_) => CastToSubPixelsError::Alignment {
            src_type_name: type_name::<&[u8]>(),
            dst_type_name: type_name::<&[Pixel]>(),
            src: bytes.as_ptr(),
            dst_align: align_of_val::<[Pixel]>(&[]),
        },
        zerocopy::ConvertError::Size(_) => CastToSubPixelsError::Size {
            src_type_name: type_name::<&[u8]>(),
            dst_type_name: type_name::<&[Pixel]>(),
            src_size: bytes.len(),
        },
    })?;

    image_rs::ImageBuffer::from_raw(dimens.width, dimens.height, buf.to_vec())
        .ok_or(FromBytesError::NotEnoughBytesForDimensions)
}