bevy_map_scatter 0.5.0

use std::sync::Arc;

use bevy::prelude::Image;
use bevy::render::render_resource::TextureFormat;
use glam::Vec2;
use map_scatter::prelude::{Texture, TextureChannel};

/// Error returned when an [`ImageTexture`] snapshot cannot be created from a Bevy [`Image`].
#[derive(Clone, Debug, PartialEq)]
pub enum ImageTextureError {
    /// The image uses a texture format that is not supported by this CPU-side adapter.
    UnsupportedFormat(TextureFormat),
    /// The image has no CPU-side pixel data.
    MissingImageData,
    /// The image has invalid dimensions for CPU texture sampling.
    InvalidDimensions {
        width: u32,
        height: u32,
        depth_or_array_layers: u32,
    },
    /// The texture domain extent must be positive and finite on both axes.
    InvalidDomainExtent { domain_extent: Vec2 },
    /// The image data length does not match the expected length for the supported format.
    UnexpectedPixelBufferLength { expected: usize, actual: usize },
}

impl std::fmt::Display for ImageTextureError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::UnsupportedFormat(format) => {
                write!(f, "unsupported image texture format: {format:?}")
            }
            Self::MissingImageData => write!(f, "image has no CPU-side pixel data"),
            Self::InvalidDimensions {
                width,
                height,
                depth_or_array_layers,
            } => write!(
                f,
                "invalid image dimensions: {width}x{height}x{depth_or_array_layers}"
            ),
            Self::InvalidDomainExtent { domain_extent } => {
                write!(f, "invalid image texture domain extent: {domain_extent:?}")
            }
            Self::UnexpectedPixelBufferLength { expected, actual } => write!(
                f,
                "unexpected image pixel buffer length: expected {expected} bytes, got {actual}"
            ),
        }
    }
}

impl std::error::Error for ImageTextureError {}

/// CPU-side adapter that snapshots a Bevy [`Image`] and implements [`Texture`].
/// This copies the pixel data into memory. Re-create the [`ImageTexture`] when the source
///   [`Image`] changes.
#[derive(Debug)]
pub struct ImageTexture {
    domain_extent: Vec2,
    domain_center: Vec2,
    format: TextureFormat,
    pixels: Arc<Vec<u8>>,
    width: u32,
    height: u32,
}

impl ImageTexture {
    /// Creates an [`ImageTexture`] snapshot from a Bevy [`Image`] and maps it to a specified domain extent.
    pub fn from_image(image: &Image, domain_extent: Vec2) -> Option<Self> {
        Self::try_from_image(image, domain_extent).ok()
    }

    /// Creates an [`ImageTexture`] snapshot from a Bevy [`Image`] and maps it to a specified domain.
    pub fn from_image_with_center(
        image: &Image,
        domain_extent: Vec2,
        domain_center: Vec2,
    ) -> Option<Self> {
        Self::try_from_image_with_center(image, domain_extent, domain_center).ok()
    }

    /// Tries to create an [`ImageTexture`] snapshot from a Bevy [`Image`].
    pub fn try_from_image(image: &Image, domain_extent: Vec2) -> Result<Self, ImageTextureError> {
        Self::try_from_image_with_center(image, domain_extent, Vec2::ZERO)
    }

    /// Tries to create an [`ImageTexture`] snapshot from a Bevy [`Image`] and maps it to a specified domain.
    pub fn try_from_image_with_center(
        image: &Image,
        domain_extent: Vec2,
        domain_center: Vec2,
    ) -> Result<Self, ImageTextureError> {
        let format = image.texture_descriptor.format;
        let bpp =
            bytes_per_pixel_for(format).ok_or(ImageTextureError::UnsupportedFormat(format))?;

        if !domain_extent.x.is_finite()
            || !domain_extent.y.is_finite()
            || domain_extent.x <= 0.0
            || domain_extent.y <= 0.0
        {
            return Err(ImageTextureError::InvalidDomainExtent { domain_extent });
        }

        let width = image.texture_descriptor.size.width;
        let height = image.texture_descriptor.size.height;
        let depth_or_array_layers = image.texture_descriptor.size.depth_or_array_layers;
        if width == 0 || height == 0 || depth_or_array_layers == 0 {
            return Err(ImageTextureError::InvalidDimensions {
                width,
                height,
                depth_or_array_layers,
            });
        }

        let pixels = image
            .data
            .clone()
            .ok_or(ImageTextureError::MissingImageData)?;
        let expected = expected_buffer_len(width, height, depth_or_array_layers, bpp)?;
        let actual = pixels.len();
        if actual != expected {
            return Err(ImageTextureError::UnexpectedPixelBufferLength { expected, actual });
        }

        Ok(Self {
            domain_extent,
            domain_center,
            format,
            pixels: Arc::new(pixels),
            width,
            height,
        })
    }

    #[inline]
    fn bytes_per_pixel(&self) -> usize {
        match self.format {
            TextureFormat::R8Unorm => 1,
            TextureFormat::Rgba8Unorm
            | TextureFormat::Rgba8UnormSrgb
            | TextureFormat::Bgra8Unorm
            | TextureFormat::Bgra8UnormSrgb => 4,
            _ => 0,
        }
    }

    #[inline]
    fn channel_offset(&self, channel: TextureChannel) -> Option<usize> {
        match self.format {
            TextureFormat::R8Unorm => match channel {
                TextureChannel::R => Some(0),
                _ => None,
            },
            TextureFormat::Rgba8Unorm | TextureFormat::Rgba8UnormSrgb => match channel {
                TextureChannel::R => Some(0),
                TextureChannel::G => Some(1),
                TextureChannel::B => Some(2),
                TextureChannel::A => Some(3),
            },
            TextureFormat::Bgra8Unorm | TextureFormat::Bgra8UnormSrgb => match channel {
                TextureChannel::B => Some(0),
                TextureChannel::G => Some(1),
                TextureChannel::R => Some(2),
                TextureChannel::A => Some(3),
            },
            _ => None,
        }
    }
}

fn bytes_per_pixel_for(format: TextureFormat) -> Option<usize> {
    match format {
        TextureFormat::R8Unorm => Some(1),
        TextureFormat::Rgba8Unorm
        | TextureFormat::Rgba8UnormSrgb
        | TextureFormat::Bgra8Unorm
        | TextureFormat::Bgra8UnormSrgb => Some(4),
        _ => None,
    }
}

fn expected_buffer_len(
    width: u32,
    height: u32,
    depth_or_array_layers: u32,
    bytes_per_pixel: usize,
) -> Result<usize, ImageTextureError> {
    let texels = (width as usize)
        .checked_mul(height as usize)
        .and_then(|n| n.checked_mul(depth_or_array_layers as usize));
    texels.and_then(|n| n.checked_mul(bytes_per_pixel)).ok_or(
        ImageTextureError::UnexpectedPixelBufferLength {
            expected: usize::MAX,
            actual: 0,
        },
    )
}

impl Texture for ImageTexture {
    fn sample(&self, channel: TextureChannel, p: Vec2) -> f32 {
        let bpp = self.bytes_per_pixel();
        if bpp == 0 {
            return 0.0;
        }

        // Map world/domain coordinates to image texels using a centered domain, like overlays.
        let (w, h) = (self.width, self.height);
        if w == 0 || h == 0 {
            return 0.0;
        }
        let (dw, dh) = (self.domain_extent.x, self.domain_extent.y);
        if dw == 0.0 || dh == 0.0 {
            return 0.0;
        }
        let local = p - self.domain_center;
        let u = ((local.x / dw) + 0.5).clamp(0.0, 1.0);
        let v = ((local.y / dh) + 0.5).clamp(0.0, 1.0);
        let x = ((u * w as f32) as u32).min(w.saturating_sub(1));
        let y = ((v * h as f32) as u32).min(h.saturating_sub(1));

        let idx = (y as usize) * (self.width as usize) + (x as usize);
        let base = idx * bpp;

        let Some(co) = self.channel_offset(channel) else {
            return 0.0;
        };

        let byte = self.pixels.get(base + co).copied().unwrap_or(0);
        (byte as f32) / 255.0
    }
}

#[cfg(test)]
mod tests {
    use bevy::asset::RenderAssetUsages;
    use bevy::render::render_resource::{Extent3d, TextureDimension};

    use super::*;

    fn image(width: u32, height: u32, format: TextureFormat, data: Vec<u8>) -> Image {
        Image::new(
            Extent3d {
                width,
                height,
                depth_or_array_layers: 1,
            },
            TextureDimension::D2,
            data,
            format,
            RenderAssetUsages::default(),
        )
    }

    fn uninit_image(width: u32, height: u32, format: TextureFormat) -> Image {
        Image::new_uninit(
            Extent3d {
                width,
                height,
                depth_or_array_layers: 1,
            },
            TextureDimension::D2,
            format,
            RenderAssetUsages::default(),
        )
    }

    #[test]
    fn creates_supported_rgba_image_snapshot() {
        let image = image(1, 1, TextureFormat::Rgba8Unorm, vec![64, 128, 192, 255]);

        let texture =
            ImageTexture::try_from_image(&image, Vec2::splat(1.0)).expect("valid image texture");

        assert_eq!(texture.sample(TextureChannel::R, Vec2::ZERO), 64.0 / 255.0);
        assert_eq!(texture.sample(TextureChannel::G, Vec2::ZERO), 128.0 / 255.0);
        assert_eq!(texture.sample(TextureChannel::B, Vec2::ZERO), 192.0 / 255.0);
        assert_eq!(texture.sample(TextureChannel::A, Vec2::ZERO), 1.0);
    }

    #[test]
    fn rejects_unsupported_format() {
        let image = image(1, 1, TextureFormat::R16Float, vec![0, 0]);

        let err = ImageTexture::try_from_image(&image, Vec2::splat(1.0)).unwrap_err();

        assert_eq!(
            err,
            ImageTextureError::UnsupportedFormat(TextureFormat::R16Float)
        );
        assert!(ImageTexture::from_image(&image, Vec2::splat(1.0)).is_none());
    }

    #[test]
    fn rejects_missing_image_data_without_panicking() {
        let image = uninit_image(1, 1, TextureFormat::Rgba8Unorm);

        let err = ImageTexture::try_from_image(&image, Vec2::splat(1.0)).unwrap_err();

        assert_eq!(err, ImageTextureError::MissingImageData);
        assert!(
            ImageTexture::from_image_with_center(&image, Vec2::splat(1.0), Vec2::ZERO).is_none()
        );
    }

    #[test]
    fn rejects_zero_image_size() {
        let image = image(0, 1, TextureFormat::R8Unorm, Vec::new());

        let err = ImageTexture::try_from_image(&image, Vec2::splat(1.0)).unwrap_err();

        assert_eq!(
            err,
            ImageTextureError::InvalidDimensions {
                width: 0,
                height: 1,
                depth_or_array_layers: 1,
            }
        );
    }

    #[test]
    fn rejects_zero_domain_extent() {
        let image = image(1, 1, TextureFormat::R8Unorm, vec![255]);

        let err = ImageTexture::try_from_image(&image, Vec2::new(0.0, 1.0)).unwrap_err();

        assert_eq!(
            err,
            ImageTextureError::InvalidDomainExtent {
                domain_extent: Vec2::new(0.0, 1.0),
            }
        );
    }

    #[test]
    fn rejects_unexpected_pixel_buffer_length() {
        let mut image = image(1, 1, TextureFormat::Rgba8Unorm, vec![1, 2, 3, 4]);
        image.data = Some(vec![1, 2, 3]);

        let err = ImageTexture::try_from_image(&image, Vec2::splat(1.0)).unwrap_err();

        assert_eq!(
            err,
            ImageTextureError::UnexpectedPixelBufferLength {
                expected: 4,
                actual: 3,
            }
        );
    }

    #[test]
    fn r8_samples_red_channel_only() {
        let image = image(1, 1, TextureFormat::R8Unorm, vec![127]);
        let texture = ImageTexture::try_from_image(&image, Vec2::splat(1.0)).unwrap();

        assert_eq!(texture.sample(TextureChannel::R, Vec2::ZERO), 127.0 / 255.0);
        assert_eq!(texture.sample(TextureChannel::G, Vec2::ZERO), 0.0);
        assert_eq!(texture.sample(TextureChannel::B, Vec2::ZERO), 0.0);
        assert_eq!(texture.sample(TextureChannel::A, Vec2::ZERO), 0.0);
    }

    #[test]
    fn rgba_samples_channel_offsets() {
        let image = image(1, 1, TextureFormat::Rgba8Unorm, vec![10, 20, 30, 40]);
        let texture = ImageTexture::try_from_image(&image, Vec2::splat(1.0)).unwrap();

        assert_eq!(texture.sample(TextureChannel::R, Vec2::ZERO), 10.0 / 255.0);
        assert_eq!(texture.sample(TextureChannel::G, Vec2::ZERO), 20.0 / 255.0);
        assert_eq!(texture.sample(TextureChannel::B, Vec2::ZERO), 30.0 / 255.0);
        assert_eq!(texture.sample(TextureChannel::A, Vec2::ZERO), 40.0 / 255.0);
    }

    #[test]
    fn bgra_samples_channel_offsets() {
        let image = image(1, 1, TextureFormat::Bgra8Unorm, vec![10, 20, 30, 40]);
        let texture = ImageTexture::try_from_image(&image, Vec2::splat(1.0)).unwrap();

        assert_eq!(texture.sample(TextureChannel::R, Vec2::ZERO), 30.0 / 255.0);
        assert_eq!(texture.sample(TextureChannel::G, Vec2::ZERO), 20.0 / 255.0);
        assert_eq!(texture.sample(TextureChannel::B, Vec2::ZERO), 10.0 / 255.0);
        assert_eq!(texture.sample(TextureChannel::A, Vec2::ZERO), 40.0 / 255.0);
    }
}