mirage-engine 0.1.1

use std::io::Cursor;

use image::{ImageError, ImageReader, Limits};

use crate::Error;
use crate::math::UVec2;

/// Largest a loaded texture may be, across and down: what every target
/// Mirage draws to binds.
const MAX_SIZE: u32 = 8192;

/// The texture a mesh slot is sampled from: `8-bit` RGBA, sRGB-encoded, row
/// by row from the top left.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct TextureData {
    size: UVec2,
    pixels: Vec<u8>,
    pixelated: bool,
}

impl TextureData {
    /// A `size`-sized texture over `pixels`, `4` bytes per pixel.
    ///
    /// The length must match `size`; checked only in debug builds.
    pub fn rgba8(size: UVec2, pixels: Vec<u8>) -> Self {
        debug_assert_eq!(
            pixels.len() as u64,
            4 * u64::from(size.x) * u64::from(size.y),
            "a {}x{} texture needs four bytes per pixel",
            size.x,
            size.y
        );

        Self {
            size,
            pixels,
            pixelated: false,
        }
    }

    /// Marks this texture sampled from the nearest texel, so that its pixels
    /// stay pixels however large it is drawn; the default blends between them.
    #[must_use]
    pub fn pixelated(mut self) -> Self {
        self.pixelated = true;
        self
    }

    /// Pixels across and down; `(0, 0)` when there are none.
    pub fn size(&self) -> UVec2 {
        self.size
    }

    /// Whether the texture has pixels to upload.
    pub(crate) fn drawn(&self) -> bool {
        self.size.x > 0 && self.size.y > 0
    }

    /// The pixels, `4` bytes each, row by row from the top left.
    pub fn pixels(&self) -> &[u8] {
        &self.pixels
    }
}

/// The relief a mesh slot reads: `8-bit` RGBA holding a normal and, where
/// the constructor declares one, a depth per texel in place of color, row by
/// row from the top left.
///
/// A relief holds no sampler of its own; it is sampled the way its slot's
/// color texture is.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct ReliefData {
    map: TextureData,
    deep: bool,
}

impl ReliefData {
    /// A `size`-sized relief over `pixels`, `4` bytes per pixel: a normal in
    /// `RGB` and a depth in `A`.
    ///
    /// The length must match `size`; checked only in debug builds.
    pub fn rgba8(size: UVec2, pixels: Vec<u8>) -> Self {
        Self::held(TextureData::rgba8(size, pixels), true)
    }

    /// The same over `pixels` holding normals alone, whose alpha byte no
    /// draw reads.
    ///
    /// Required if you want a relief that lights a surface without moving
    /// its texels off the plane; a `.glb` normal texture is read as one.
    pub fn normals(size: UVec2, pixels: Vec<u8>) -> Self {
        Self::held(TextureData::rgba8(size, pixels), false)
    }

    /// Pixels across and down; `(0, 0)` when there are none.
    pub fn size(&self) -> UVec2 {
        self.map.size()
    }

    /// The relief a source loaded as a texture, whose alpha holds a depth.
    pub(crate) fn loaded(texture: TextureData) -> Self {
        Self::held(texture, true)
    }

    /// Whether the relief holds a depth beside its normals, which is what
    /// its constructor declared and never what its pixels hold.
    pub(crate) fn deep(&self) -> bool {
        self.deep
    }

    /// The pixels to upload, whose channels the shader reads as a normal and
    /// a depth per texel.
    pub(crate) fn map(&self) -> &TextureData {
        &self.map
    }

    const fn held(map: TextureData, deep: bool) -> Self {
        Self { map, deep }
    }
}

/// The shading a mesh slot reads: `8-bit` RGBA holding a texel's occlusion,
/// roughness and metallic in `R`, `G` and `B` in place of color, row by row
/// from the top left.
///
/// A shading map holds no sampler of its own; it is sampled the way its
/// slot's color texture is.
#[derive(Clone, Debug, PartialEq)]
pub struct ShadingData(TextureData);

impl ShadingData {
    /// A `size`-sized shading map over `pixels`, `4` bytes per pixel.
    ///
    /// The length must match `size`; checked only in debug builds.
    pub fn rgba8(size: UVec2, pixels: Vec<u8>) -> Self {
        Self(TextureData::rgba8(size, pixels))
    }

    /// The pixels to upload, whose channels the shader reads as an
    /// occlusion, a roughness and a metallic per texel.
    pub(crate) fn map(&self) -> &TextureData {
        &self.0
    }
}

/// The GPU side of slot textures: the samplers a slot reads through, the
/// white pixel a slot with no texture is drawn against, or its shading and
/// emissive maps read where it holds none, and the flat texel a slot with
/// no relief reads.
pub(crate) struct Textures {
    layout: wgpu::BindGroupLayout,
    blending: wgpu::Sampler,
    nearest: wgpu::Sampler,
    white: wgpu::Texture,
    flat: wgpu::Texture,
    fallback: wgpu::BindGroup,
}

impl Textures {
    pub(crate) fn new(device: &wgpu::Device, queue: &wgpu::Queue) -> Self {
        let map = |binding| wgpu::BindGroupLayoutEntry {
            binding,
            visibility: wgpu::ShaderStages::FRAGMENT,
            ty: wgpu::BindingType::Texture {
                sample_type: wgpu::TextureSampleType::Float { filterable: true },
                view_dimension: wgpu::TextureViewDimension::D2,
                multisampled: false,
            },
            count: None,
        };
        let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("mirage-engine slot texture"),
            entries: &[
                map(0),
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    visibility: wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
                    count: None,
                },
                map(2),
                map(3),
                map(4),
            ],
        });
        let blending = sampler(device, wgpu::FilterMode::Linear);
        let nearest = sampler(device, wgpu::FilterMode::Nearest);

        let white = uploaded(
            device,
            queue,
            &TextureData::rgba8(UVec2::ONE, vec![u8::MAX; 4]),
            wgpu::TextureFormat::Rgba8UnormSrgb,
        );
        // The relief a slot with none reads: a normal straight out of the
        // sprite, at no depth.
        let flat = uploaded(
            device,
            queue,
            &TextureData::rgba8(UVec2::ONE, vec![128, 128, u8::MAX, 0]),
            wgpu::TextureFormat::Rgba8Unorm,
        );
        let fallback = bindings(device, &layout, &blending, [&white, &flat, &white, &white]);
        Self {
            layout,
            blending,
            nearest,
            white,
            flat,
            fallback,
        }
    }

    pub(crate) fn layout(&self) -> &wgpu::BindGroupLayout {
        &self.layout
    }

    /// Fallback a slot with no texture samples: white, so shading is only
    /// the slot's tint.
    pub(crate) fn fallback(&self) -> &wgpu::BindGroup {
        &self.fallback
    }

    /// Uploads a slot's `color` and the maps beside it and binds them
    /// together, or `None` where none of them has pixels.
    ///
    /// The relief and the shading are uploaded raw: their channels hold
    /// normals, depths and factors, not color. The slot's color chooses the
    /// sampler, which all four read through.
    pub(crate) fn bind(
        &self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        color: Option<&TextureData>,
        relief: Option<&ReliefData>,
        shading: Option<&ShadingData>,
        emissive: Option<&TextureData>,
    ) -> Option<wgpu::BindGroup> {
        let color = color.filter(|data| data.drawn());
        let relief = relief.map(ReliefData::map).filter(|data| data.drawn());
        let shading = shading.map(ShadingData::map).filter(|data| data.drawn());
        let emissive = emissive.filter(|data| data.drawn());
        if [color, relief, shading, emissive]
            .iter()
            .all(Option::is_none)
        {
            return None;
        }
        let sampler = match color.is_some_and(|data| data.pixelated) {
            true => &self.nearest,
            false => &self.blending,
        };
        let paint =
            |data: &TextureData| uploaded(device, queue, data, wgpu::TextureFormat::Rgba8UnormSrgb);
        let raw =
            |data: &TextureData| uploaded(device, queue, data, wgpu::TextureFormat::Rgba8Unorm);
        let (base, raised) = (color.map(paint), relief.map(raw));
        let (scaled, cast) = (shading.map(raw), emissive.map(paint));

        Some(bindings(
            device,
            &self.layout,
            sampler,
            [
                base.as_ref().unwrap_or(&self.white),
                raised.as_ref().unwrap_or(&self.flat),
                scaled.as_ref().unwrap_or(&self.white),
                cast.as_ref().unwrap_or(&self.white),
            ],
        ))
    }
}

/// The texture `bytes` hold, up to [`MAX_SIZE`] pixels a side.
///
/// The size is read before any pixels are held, so a source that declares
/// more than a target binds fails without holding the memory it declared.
pub(crate) fn decode(bytes: &[u8]) -> Result<TextureData, Error> {
    let mut reader = ImageReader::new(Cursor::new(bytes))
        .with_guessed_format()
        .map_err(|error| Error::msg(format!("did not decode: {error}")))?;
    reader.limits(bounded());

    let decoded = reader.decode().map_err(refused)?.into_rgba8();
    let size = UVec2::new(decoded.width(), decoded.height());

    Ok(TextureData::rgba8(size, decoded.into_raw()))
}

/// Largest a source may declare itself before the decoder stops reading it.
fn bounded() -> Limits {
    let mut limits = Limits::default();
    limits.max_image_width = Some(MAX_SIZE);
    limits.max_image_height = Some(MAX_SIZE);

    limits
}

/// The error for a source that did not decode. Where the size cap stopped
/// it, the error states the cap.
fn refused(error: ImageError) -> Error {
    match error {
        ImageError::Limits(_) => Error::msg(format!(
            "is larger than the {MAX_SIZE} pixels a side Mirage draws"
        )),
        error => Error::msg(format!("did not decode: {error}")),
    }
}

/// Sample filter between a texture's texels, the one thing
/// [`TextureData::pixelated`] changes.
fn sampler(device: &wgpu::Device, filter: wgpu::FilterMode) -> wgpu::Sampler {
    device.create_sampler(&wgpu::SamplerDescriptor {
        label: Some("mirage-engine slot texture"),
        address_mode_u: wgpu::AddressMode::Repeat,
        address_mode_v: wgpu::AddressMode::Repeat,
        address_mode_w: wgpu::AddressMode::Repeat,
        mag_filter: filter,
        min_filter: filter,
        ..Default::default()
    })
}

/// Binds `sampler` and the slot's color, relief, shading and emissive maps,
/// in that order.
fn bindings(
    device: &wgpu::Device,
    layout: &wgpu::BindGroupLayout,
    sampler: &wgpu::Sampler,
    maps: [&wgpu::Texture; 4],
) -> wgpu::BindGroup {
    let [color, relief, shading, emissive] = maps.map(|map| map.create_view(&Default::default()));
    let map = |binding, view| wgpu::BindGroupEntry {
        binding,
        resource: wgpu::BindingResource::TextureView(view),
    };
    device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("mirage-engine slot texture"),
        layout,
        entries: &[
            map(0, &color),
            wgpu::BindGroupEntry {
                binding: 1,
                resource: wgpu::BindingResource::Sampler(sampler),
            },
            map(2, &relief),
            map(3, &shading),
            map(4, &emissive),
        ],
    })
}

/// Uploads `data` as a texture of `format`.
fn uploaded(
    device: &wgpu::Device,
    queue: &wgpu::Queue,
    data: &TextureData,
    format: wgpu::TextureFormat,
) -> wgpu::Texture {
    let extent = wgpu::Extent3d {
        width: data.size().x,
        height: data.size().y,
        depth_or_array_layers: 1,
    };
    let texture = device.create_texture(&wgpu::TextureDescriptor {
        label: Some("mirage-engine slot texture"),
        size: extent,
        mip_level_count: 1,
        sample_count: 1,
        dimension: wgpu::TextureDimension::D2,
        format,
        usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
        view_formats: &[],
    });
    queue.write_texture(
        wgpu::TexelCopyTextureInfo {
            texture: &texture,
            mip_level: 0,
            origin: wgpu::Origin3d::ZERO,
            aspect: wgpu::TextureAspect::All,
        },
        data.pixels(),
        wgpu::TexelCopyBufferLayout {
            offset: 0,
            bytes_per_row: Some(4 * data.size().x),
            rows_per_image: Some(data.size().y),
        },
        extent,
    );

    texture
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::assets::IMP;

    /// A source of `width` by `height` pixels, as a file would hold it.
    fn png(width: u32, height: u32) -> Vec<u8> {
        let mut out = Vec::new();
        image::RgbaImage::new(width, height)
            .write_to(&mut Cursor::new(&mut out), image::ImageFormat::Png)
            .expect("the fixture encodes");

        out
    }

    #[test]
    fn a_source_larger_than_a_target_binds_does_not_decode() {
        let error = decode(&png(MAX_SIZE + 1, 1)).expect_err("no target binds that");

        assert_eq!(
            error.to_string(),
            format!("is larger than the {MAX_SIZE} pixels a side Mirage draws")
        );
        assert_eq!(
            decode(&png(MAX_SIZE, 1))
                .expect("the cap itself is drawn")
                .size(),
            UVec2::new(MAX_SIZE, 1),
        );
    }

    #[test]
    fn a_source_cut_off_anywhere_reads_as_itself_or_as_an_error() {
        let whole = decode(IMP).expect("the fixture decodes").size();

        for at in 0..IMP.len() {
            if let Ok(cut) = decode(&IMP[..at]) {
                assert_eq!(cut.size(), whole, "a cut at {at}");
            }
        }
    }
}