kornia-io 0.2.0

use crate::{
    conv_utils::{convert_buf_u16_u8, convert_buf_u8_u16_into_slice},
    error::IoError,
    limits::{alloc_image, check_image_dimensions, try_alloc_zeroed},
};
use kornia_image::{
    color_spaces::{Gray16, Gray8, Rgb16, Rgb8, Rgba16, Rgba8},
    Image, ImageLayout, ImageSize, PixelFormat,
};
use png::{BitDepth, ColorType, Decoder, DeflateCompression, Encoder};
use std::{
    fs,
    fs::File,
    io::{BufReader, Cursor, Write},
    path::Path,
};

/// Read a PNG image as grayscale (Gray8).
///
/// # Arguments
///
/// * `file_path` - The path to the PNG file.
///
/// # Returns
///
/// A grayscale image (Gray8).
pub fn read_image_png_mono8(file_path: impl AsRef<Path>) -> Result<Gray8, IoError> {
    Ok(Gray8(read_png_u8(file_path)?))
}

/// Read a PNG image as RGB8.
///
/// # Arguments
///
/// * `file_path` - The path to the PNG file.
///
/// # Returns
///
/// An RGB8 typed image.
pub fn read_image_png_rgb8(file_path: impl AsRef<Path>) -> Result<Rgb8, IoError> {
    Ok(Rgb8(read_png_u8(file_path)?))
}

/// Read a PNG image as RGBA8.
///
/// # Arguments
///
/// * `file_path` - The path to the PNG file.
///
/// # Returns
///
/// An RGBA8 typed image.
pub fn read_image_png_rgba8(file_path: impl AsRef<Path>) -> Result<Rgba8, IoError> {
    Ok(Rgba8(read_png_u8(file_path)?))
}

/// Read a PNG image as RGB16.
///
/// # Arguments
///
/// * `file_path` - The path to the PNG file.
///
/// # Returns
///
/// An RGB16 typed image.
pub fn read_image_png_rgb16(file_path: impl AsRef<Path>) -> Result<Rgb16, IoError> {
    Ok(Rgb16(read_png_u16(file_path)?))
}

/// Read a PNG image as RGBA16.
///
/// # Arguments
///
/// * `file_path` - The path to the PNG file.
///
/// # Returns
///
/// An RGBA16 typed image.
pub fn read_image_png_rgba16(file_path: impl AsRef<Path>) -> Result<Rgba16, IoError> {
    Ok(Rgba16(read_png_u16(file_path)?))
}

/// Read a PNG image as grayscale (Gray16).
///
/// # Arguments
///
/// * `file_path` - The path to the PNG file.
///
/// # Returns
///
/// A Gray16 typed image.
pub fn read_image_png_mono16(file_path: impl AsRef<Path>) -> Result<Gray16, IoError> {
    Ok(Gray16(read_png_u16(file_path)?))
}

/// Decodes a PNG image with as grayscale (Gray8) from Raw Bytes.
///
/// # Arguments
///
/// - `src` - Raw bytes of the png file
/// - `dst` - A mutable reference to your `Gray8` image
pub fn decode_image_png_mono8(src: &[u8], dst: &mut Gray8) -> Result<(), IoError> {
    let size = dst.size();
    decode_png_impl(src, dst.as_slice_mut(), size, 1, BitDepth::Eight)
}

/// Decodes a PNG image with a three channel (rgb8) from Raw Bytes.
///
/// # Arguments
///
/// - `src` - Raw bytes of the png file
/// - `dst` - A mutable reference to your `Rgb8` image
pub fn decode_image_png_rgb8(src: &[u8], dst: &mut Rgb8) -> Result<(), IoError> {
    let size = dst.size();
    decode_png_impl(src, dst.as_slice_mut(), size, 3, BitDepth::Eight)
}

/// Decodes a PNG image with a four channel (rgba8) from Raw Bytes.
///
/// # Arguments
///
/// - `src` - Raw bytes of the png file
/// - `dst` - A mutable reference to your `Rgba8` image
pub fn decode_image_png_rgba8(src: &[u8], dst: &mut Rgba8) -> Result<(), IoError> {
    let size = dst.size();
    decode_png_impl(src, dst.as_slice_mut(), size, 4, BitDepth::Eight)
}

/// Decodes a PNG (16 Bit) image as grayscale (Gray16) from Raw Bytes.
///
/// # Arguments
///
/// - `src` - Raw bytes of the png file
/// - `dst` - A mutable reference to your `Gray16` image
pub fn decode_image_png_mono16(src: &[u8], dst: &mut Gray16) -> Result<(), IoError> {
    let mut image_u8 = convert_buf_u16_u8(dst.as_slice());
    decode_png_impl(
        src,
        image_u8.as_mut_slice(),
        dst.size(),
        1,
        BitDepth::Sixteen,
    )?;
    convert_buf_u8_u16_into_slice(image_u8.as_slice(), dst.as_slice_mut());
    Ok(())
}

/// Decodes a PNG (16 Bit) image with a three channel (rgb16) from Raw Bytes.
///
/// # Arguments
///
/// - `src` - Raw bytes of the png file
/// - `dst` - A mutable reference to your `Rgb16` image
pub fn decode_image_png_rgb16(src: &[u8], dst: &mut Rgb16) -> Result<(), IoError> {
    let mut image_u8 = convert_buf_u16_u8(dst.as_slice());
    decode_png_impl(
        src,
        image_u8.as_mut_slice(),
        dst.size(),
        3,
        BitDepth::Sixteen,
    )?;
    convert_buf_u8_u16_into_slice(image_u8.as_slice(), dst.as_slice_mut());
    Ok(())
}

/// Decodes a PNG (16 Bit) image with as RGBA (Rgba16) from Raw Bytes.
///
/// # Arguments
///
/// - `src` - Raw bytes of the png file
/// - `dst` - A mutable reference to your `Rgba16` image
pub fn decode_image_png_rgba16(src: &[u8], dst: &mut Rgba16) -> Result<(), IoError> {
    let mut image_u8 = convert_buf_u16_u8(dst.as_slice());
    decode_png_impl(
        src,
        image_u8.as_mut_slice(),
        dst.size(),
        4,
        BitDepth::Sixteen,
    )?;
    convert_buf_u8_u16_into_slice(image_u8.as_slice(), dst.as_slice_mut());
    Ok(())
}

/// Decodes PNG image metadata from raw bytes without decoding pixel data.
///
/// # Arguments
///
/// - `src` - Raw bytes of the PNG file
///
/// # Returns
///
/// An `ImageLayout` containing the image metadata (size, channels, pixel format).
pub fn decode_image_png_layout(src: &[u8]) -> Result<ImageLayout, IoError> {
    let cursor = Cursor::new(src);
    let decoder = Decoder::new(cursor);
    let reader = decoder
        .read_info()
        .map_err(|e| IoError::PngDecodeError(e.to_string()))?;

    let info = reader.info();
    let size = ImageSize {
        width: info.width as usize,
        height: info.height as usize,
    };
    check_image_dimensions(size.width, size.height)?;

    let channels: u8 = match info.color_type {
        ColorType::Grayscale => 1,
        ColorType::Rgb => 3,
        ColorType::Rgba => 4,
        ColorType::GrayscaleAlpha => 2,
        ColorType::Indexed => 1,
    };

    let pixel_format = match info.bit_depth {
        BitDepth::Eight => PixelFormat::U8,
        BitDepth::Sixteen => PixelFormat::U16,
        other => {
            return Err(IoError::PngDecodeError(format!(
                "Unsupported bit depth: {:?}",
                other
            )))
        }
    };

    Ok(ImageLayout::new(size, channels, pixel_format))
}

// Validates the decoded pixel format and dimensions against what the caller requested and
// returns the exact output buffer size in bytes.
fn check_png_output<R: std::io::BufRead + std::io::Seek>(
    reader: &png::Reader<R>,
    channels: usize,
    bit_depth: BitDepth,
) -> Result<usize, IoError> {
    let info = reader.info();
    check_image_dimensions(info.width as usize, info.height as usize)?;

    let (out_color, out_depth) = reader.output_color_type();
    if out_color.samples() != channels || out_depth != bit_depth {
        return Err(IoError::FormatMismatch(format!(
            "PNG is {out_color:?} {out_depth:?}, expected {channels} channel(s) at {bit_depth:?}"
        )));
    }

    reader
        .output_buffer_size()
        .ok_or_else(|| IoError::PngDecodeError("PNG output buffer size overflowed".into()))
}

// Opens a PNG file after validating that it exists and has a `.png` extension.
fn open_png(file_path: impl AsRef<Path>) -> Result<png::Reader<BufReader<File>>, IoError> {
    let file_path = file_path.as_ref();
    if !file_path.exists() {
        return Err(IoError::FileDoesNotExist(file_path.to_path_buf()));
    }
    if file_path.extension().is_none_or(|ext| ext != "png") {
        return Err(IoError::InvalidFileExtension(file_path.to_path_buf()));
    }

    let file = fs::File::open(file_path)?;
    Decoder::new(BufReader::new(file))
        .read_info()
        .map_err(|e| IoError::PngDecodeError(e.to_string()))
}

// Validates the output format of `reader` against a `C`-channel image of `T` samples and
// allocates it. Returns the image and the decoder's output size in bytes.
fn alloc_png_output<R: std::io::BufRead + std::io::Seek, T: Clone + Default, const C: usize>(
    reader: &png::Reader<R>,
    bit_depth: BitDepth,
) -> Result<(Image<T, C>, usize), IoError> {
    let buffer_size = check_png_output(reader, C, bit_depth)?;
    let (width, height) = reader.info().size();
    let img = alloc_image::<T, C>(ImageSize {
        width: width as usize,
        height: height as usize,
    })?;
    let expected = std::mem::size_of_val(img.as_slice());
    if buffer_size != expected {
        return Err(IoError::InvalidBufferSize(buffer_size, expected));
    }
    Ok((img, buffer_size))
}

// Reads an 8-bit PNG file straight into a newly allocated image.
fn read_png_u8<const C: usize>(file_path: impl AsRef<Path>) -> Result<Image<u8, C>, IoError> {
    let mut reader = open_png(file_path)?;
    let (mut img, _) = alloc_png_output::<_, u8, C>(&reader, BitDepth::Eight)?;
    reader
        .next_frame(img.as_slice_mut())
        .map_err(|e| IoError::PngDecodeError(e.to_string()))?;
    Ok(img)
}

// Reads a 16-bit PNG file (big-endian samples) into a newly allocated image.
fn read_png_u16<const C: usize>(file_path: impl AsRef<Path>) -> Result<Image<u16, C>, IoError> {
    let mut reader = open_png(file_path)?;
    let (mut img, buffer_size) = alloc_png_output::<_, u16, C>(&reader, BitDepth::Sixteen)?;
    let mut buf = try_alloc_zeroed::<u8>(buffer_size)?;
    reader
        .next_frame(&mut buf)
        .map_err(|e| IoError::PngDecodeError(e.to_string()))?;
    convert_buf_u8_u16_into_slice(&buf, img.as_slice_mut());
    Ok(img)
}

// Utility function to decode png files from raw bytes
fn decode_png_impl(
    src: &[u8],
    dst: &mut [u8],
    image_size: ImageSize,
    channels: usize,
    bit_depth: BitDepth,
) -> Result<(), IoError> {
    let cursor = Cursor::new(src);
    let mut reader = Decoder::new(cursor)
        .read_info()
        .map_err(|e| IoError::PngDecodeError(e.to_string()))?;

    let image_info = reader.info();
    if image_info.size() != (image_size.width as u32, image_size.height as u32) {
        return Err(IoError::DecodeMismatchResolution(
            image_info.height as usize,
            image_info.width as usize,
            image_size.height,
            image_size.width,
        ));
    }

    let buffer_size = check_png_output(&reader, channels, bit_depth)?;

    if dst.len() != buffer_size {
        return Err(IoError::InvalidBufferSize(dst.len(), buffer_size));
    }

    let _ = reader
        .next_frame(dst)
        .map_err(|e| IoError::PngDecodeError(e.to_string()))?;

    Ok(())
}

/// Writes the given PNG _(rgb8)_ data to the given file path.
///
/// # Arguments
///
/// - `file_path` - The path to the PNG image.
/// - `image` - The Rgb8 image to write.
pub fn write_image_png_rgb8(
    file_path: impl AsRef<Path>,
    image: &Image<u8, 3>,
) -> Result<(), IoError> {
    write_png_impl(
        file_path,
        image.as_slice(),
        image.size(),
        BitDepth::Eight,
        ColorType::Rgb,
    )
}

/// Writes the given PNG _(rgba8)_ data to the given file path.
///
/// # Arguments
///
/// - `file_path` - The path to the PNG image.
/// - `image` - The Rgba8 image to write.
pub fn write_image_png_rgba8(
    file_path: impl AsRef<Path>,
    image: &Image<u8, 4>,
) -> Result<(), IoError> {
    write_png_impl(
        file_path,
        image.as_slice(),
        image.size(),
        BitDepth::Eight,
        ColorType::Rgba,
    )
}

/// Writes the given PNG _(grayscale 8-bit)_ data to the given file path.
///
/// # Arguments
///
/// - `file_path` - The path to the PNG image.
/// - `image` - The Gray8 image to write.
pub fn write_image_png_gray8(
    file_path: impl AsRef<Path>,
    image: &Image<u8, 1>,
) -> Result<(), IoError> {
    write_png_impl(
        file_path,
        image.as_slice(),
        image.size(),
        BitDepth::Eight,
        ColorType::Grayscale,
    )
}

/// Writes the given PNG _(rgb16)_ data to the given file path.
///
/// # Arguments
///
/// - `file_path` - The path to the PNG image.
/// - `image` - The Rgb16 image to write.
pub fn write_image_png_rgb16(
    file_path: impl AsRef<Path>,
    image: &Image<u16, 3>,
) -> Result<(), IoError> {
    let image_size = image.size();
    let image_buf = convert_buf_u16_u8(image.as_slice());

    write_png_impl(
        file_path,
        &image_buf,
        image_size,
        BitDepth::Sixteen,
        ColorType::Rgb,
    )
}

/// Writes the given PNG _(rgba16)_ data to the given file path.
///
/// # Arguments
///
/// - `file_path` - The path to the PNG image.
/// - `image` - The Rgba16 image to write.
pub fn write_image_png_rgba16(
    file_path: impl AsRef<Path>,
    image: &Image<u16, 4>,
) -> Result<(), IoError> {
    let image_size = image.size();
    let image_buf = convert_buf_u16_u8(image.as_slice());

    write_png_impl(
        file_path,
        &image_buf,
        image_size,
        BitDepth::Sixteen,
        ColorType::Rgba,
    )
}

/// Writes the given PNG _(grayscale 16-bit)_ data to the given file path.
///
/// # Arguments
///
/// - `file_path` - The path to the PNG image.
/// - `image` - The Gray16 image to write.
pub fn write_image_png_gray16(
    file_path: impl AsRef<Path>,
    image: &Image<u16, 1>,
) -> Result<(), IoError> {
    let image_size = image.size();
    let image_buf = convert_buf_u16_u8(image.as_slice());

    write_png_impl(
        file_path,
        &image_buf,
        image_size,
        BitDepth::Sixteen,
        ColorType::Grayscale,
    )
}

/// Writes PNG-encoded image data into any `Write` target — used by both the
/// file-path `write_image_png_*` API and the in-memory `encode_image_png_*` API.
///
/// Callers are responsible for matching `depth`/`color_type` to the layout of
/// `image_data` (e.g. 16-bit data must be passed as big-endian byte pairs and
/// `BitDepth::Sixteen`).
/// Maps a 0..=9 compression level (zlib convention) to png-crate's
/// [`DeflateCompression`]. Level 1 routes to ``FdeflateUltraFast`` —
/// the same fast path the `png` crate uses internally for ``Compression::Fast``,
/// powered by the NEON/AVX2-accelerated `fdeflate` crate.
fn level_to_deflate(level: u8) -> DeflateCompression {
    match level {
        0 => DeflateCompression::NoCompression,
        1 => DeflateCompression::FdeflateUltraFast,
        n => DeflateCompression::Level(n.min(9)),
    }
}

fn write_png_into<W: Write>(
    writer: W,
    image_data: &[u8],
    image_size: ImageSize,
    depth: BitDepth,
    color_type: ColorType,
    compress_level: Option<u8>,
) -> Result<(), IoError> {
    let mut encoder = Encoder::new(writer, image_size.width as u32, image_size.height as u32);
    encoder.set_color(color_type);
    encoder.set_depth(depth);
    if let Some(level) = compress_level {
        encoder.set_deflate_compression(level_to_deflate(level));
    }

    let mut writer = encoder
        .write_header()
        .map_err(|e| IoError::PngEncodingError(e.to_string()))?;
    writer
        .write_image_data(image_data)
        .map_err(|e| IoError::PngEncodingError(e.to_string()))?;
    Ok(())
}

fn write_png_impl(
    file_path: impl AsRef<Path>,
    image_data: &[u8],
    image_size: ImageSize,
    depth: BitDepth,
    color_type: ColorType,
) -> Result<(), IoError> {
    let file = File::create(file_path)?;
    write_png_into(file, image_data, image_size, depth, color_type, None)
}

// In-memory encoders. Buffer is appended to (caller clears for fresh encode,
// or reuses the allocation across frames). 16-bit variants serialize `&[u16]`
// to big-endian byte pairs as required by the PNG wire format.

/// Encodes an RGB8 image as PNG bytes into `buffer`. Encoded data is
/// appended (call `buffer.clear()` first to reuse the buffer fresh).
///
/// `compress_level` follows the zlib convention: 0 = no compression
/// (fastest), 1 = ``FdeflateUltraFast`` (the NEON/AVX2 fast path; ~3×
/// faster than the default at moderately worse compression), 2..9 =
/// zlib levels (default is 6 / "balanced"). ``None`` keeps the
/// `png` crate default.
pub fn encode_image_png_rgb8(
    image: &Image<u8, 3>,
    buffer: &mut Vec<u8>,
    compress_level: Option<u8>,
) -> Result<(), IoError> {
    buffer.reserve(image.as_slice().len() / 2);
    write_png_into(
        buffer,
        image.as_slice(),
        image.size(),
        BitDepth::Eight,
        ColorType::Rgb,
        compress_level,
    )
}

/// Encodes an RGBA8 image as PNG bytes into `buffer`. See
/// [`encode_image_png_rgb8`] for `compress_level` semantics.
pub fn encode_image_png_rgba8(
    image: &Image<u8, 4>,
    buffer: &mut Vec<u8>,
    compress_level: Option<u8>,
) -> Result<(), IoError> {
    buffer.reserve(image.as_slice().len() / 2);
    write_png_into(
        buffer,
        image.as_slice(),
        image.size(),
        BitDepth::Eight,
        ColorType::Rgba,
        compress_level,
    )
}

/// Encodes a grayscale 8-bit image as PNG bytes into `buffer`. See
/// [`encode_image_png_rgb8`] for `compress_level` semantics.
pub fn encode_image_png_gray8(
    image: &Image<u8, 1>,
    buffer: &mut Vec<u8>,
    compress_level: Option<u8>,
) -> Result<(), IoError> {
    buffer.reserve(image.as_slice().len() / 2);
    write_png_into(
        buffer,
        image.as_slice(),
        image.size(),
        BitDepth::Eight,
        ColorType::Grayscale,
        compress_level,
    )
}

/// Encodes an RGB16 image as PNG bytes into `buffer`. See
/// [`encode_image_png_rgb8`] for `compress_level` semantics.
pub fn encode_image_png_rgb16(
    image: &Image<u16, 3>,
    buffer: &mut Vec<u8>,
    compress_level: Option<u8>,
) -> Result<(), IoError> {
    let image_size = image.size();
    let image_buf = convert_buf_u16_u8(image.as_slice());
    buffer.reserve(image_buf.len() / 2);
    write_png_into(
        buffer,
        &image_buf,
        image_size,
        BitDepth::Sixteen,
        ColorType::Rgb,
        compress_level,
    )
}

/// Encodes an RGBA16 image as PNG bytes into `buffer`. See
/// [`encode_image_png_rgb8`] for `compress_level` semantics.
pub fn encode_image_png_rgba16(
    image: &Image<u16, 4>,
    buffer: &mut Vec<u8>,
    compress_level: Option<u8>,
) -> Result<(), IoError> {
    let image_size = image.size();
    let image_buf = convert_buf_u16_u8(image.as_slice());
    buffer.reserve(image_buf.len() / 2);
    write_png_into(
        buffer,
        &image_buf,
        image_size,
        BitDepth::Sixteen,
        ColorType::Rgba,
        compress_level,
    )
}

/// Encodes a grayscale 16-bit image as PNG bytes into `buffer`. See
/// [`encode_image_png_rgb8`] for `compress_level` semantics.
pub fn encode_image_png_gray16(
    image: &Image<u16, 1>,
    buffer: &mut Vec<u8>,
    compress_level: Option<u8>,
) -> Result<(), IoError> {
    encode_image_png_gray16_slice(image.as_slice(), image.size(), buffer, compress_level)
}

/// Encodes a borrowed 16-bit grayscale pixel slice (`width * height` u16, row-major) into a PNG.
///
/// Zero-copy of the caller's buffer: the u16 values are borrowed and packed straight to
/// big-endian PNG bytes, so no owning [`Image`] has to be built first.
/// [`encode_image_png_gray16`] delegates here.
///
/// # Arguments
///
/// * `pixels` - The 16-bit grayscale values, row-major; `pixels.len()` must equal `width * height`.
/// * `image_size` - The image dimensions.
/// * `buffer` - The output buffer the PNG bytes are appended to.
/// * `compress_level` - Optional zlib compression level.
pub fn encode_image_png_gray16_slice(
    pixels: &[u16],
    image_size: ImageSize,
    buffer: &mut Vec<u8>,
    compress_level: Option<u8>,
) -> Result<(), IoError> {
    let expected = image_size.width * image_size.height;
    if pixels.len() != expected {
        return Err(IoError::InvalidBufferSize(pixels.len(), expected));
    }
    let image_buf = convert_buf_u16_u8(pixels);
    buffer.reserve(image_buf.len() / 2);
    write_png_into(
        buffer,
        &image_buf,
        image_size,
        BitDepth::Sixteen,
        ColorType::Grayscale,
        compress_level,
    )
}

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

    fn crc32(data: &[u8]) -> u32 {
        let mut crc = 0xFFFF_FFFFu32;
        for &b in data {
            crc ^= b as u32;
            for _ in 0..8 {
                crc = if crc & 1 != 0 {
                    (crc >> 1) ^ 0xEDB8_8320
                } else {
                    crc >> 1
                };
            }
        }
        !crc
    }

    fn chunk(out: &mut Vec<u8>, ty: &[u8; 4], data: &[u8]) {
        out.extend_from_slice(&(data.len() as u32).to_be_bytes());
        let start = out.len();
        out.extend_from_slice(ty);
        out.extend_from_slice(data);
        let crc = crc32(&out[start..]);
        out.extend_from_slice(&crc.to_be_bytes());
    }

    // A tiny PNG whose header declares a huge 8-bit grayscale image.
    fn png_bomb(width: u32, height: u32) -> Vec<u8> {
        let mut out = b"\x89PNG\r\n\x1a\n".to_vec();
        let mut ihdr = Vec::new();
        ihdr.extend_from_slice(&width.to_be_bytes());
        ihdr.extend_from_slice(&height.to_be_bytes());
        ihdr.extend_from_slice(&[8, 0, 0, 0, 0]);
        chunk(&mut out, b"IHDR", &ihdr);
        chunk(&mut out, b"IDAT", &[]);
        chunk(&mut out, b"IEND", &[]);
        out
    }

    #[test]
    fn rejects_decompression_bomb() -> Result<(), Box<dyn std::error::Error>> {
        let bomb = png_bomb(1_000_000, (1 << 31) - 1);
        assert!(matches!(
            decode_image_png_layout(&bomb),
            Err(IoError::ImageTooLarge { .. })
        ));

        let dir = tempfile::tempdir()?;
        let path = dir.path().join("bomb.png");
        std::fs::write(&path, &bomb)?;
        assert!(matches!(
            read_image_png_mono8(&path),
            Err(IoError::ImageTooLarge { .. })
        ));
        Ok(())
    }

    #[test]
    fn rejects_mismatched_pixel_format() -> Result<(), Box<dyn std::error::Error>> {
        let size = ImageSize {
            width: 4,
            height: 4,
        };
        let gray = Gray8::from_size_val(size, 200)?;
        let mut encoded = Vec::new();
        encode_image_png_gray8(&gray, &mut encoded, None)?;

        let mut rgb = Rgb8::from_size_val(size, 0)?;
        assert!(matches!(
            decode_image_png_rgb8(&encoded, &mut rgb),
            Err(IoError::FormatMismatch(_))
        ));
        let mut mono16 = Gray16::from_size_val(size, 0)?;
        assert!(matches!(
            decode_image_png_mono16(&encoded, &mut mono16),
            Err(IoError::FormatMismatch(_))
        ));
        Ok(())
    }

    #[test]
    fn encode_gray16_slice_matches_owning_and_validates_len() -> Result<(), IoError> {
        let size = ImageSize {
            width: 40,
            height: 30,
        };
        let pixels: Vec<u16> = (0..(size.width * size.height) as u32)
            .map(|i| (i % 9000) as u16)
            .collect();
        // The zero-copy slice path and the owning path must produce identical PNG bytes.
        let img = Image::<u16, 1>::from_size_slice(size, &pixels)?;
        let mut via_slice = Vec::new();
        let mut via_owning = Vec::new();
        encode_image_png_gray16_slice(&pixels, size, &mut via_slice, None)?;
        encode_image_png_gray16(&img, &mut via_owning, None)?;
        assert_eq!(via_slice, via_owning);
        // A wrong-length slice is rejected.
        let mut sink = Vec::new();
        assert!(matches!(
            encode_image_png_gray16_slice(&pixels[..pixels.len() - 1], size, &mut sink, None),
            Err(IoError::InvalidBufferSize(..))
        ));
        Ok(())
    }
    use crate::error::IoError;
    use std::fs::{create_dir_all, read};

    #[test]
    fn read_png_mono8() -> Result<(), IoError> {
        let image = read_image_png_mono8("../../tests/data/dog.png")?;
        assert_eq!(image.size().width, 258);
        assert_eq!(image.size().height, 195);
        Ok(())
    }

    #[test]
    fn read_write_png_rgb8() -> Result<(), IoError> {
        let tmp_dir = tempfile::tempdir()?;
        create_dir_all(tmp_dir.path())?;

        let file_path = tmp_dir.path().join("dog-rgb8.png");
        let image_data = read_image_png_rgb8("../../tests/data/dog-rgb8.png")?;
        write_image_png_rgb8(&file_path, &image_data)?;

        let image_data_back = read_image_png_rgb8(&file_path)?;
        assert!(file_path.exists(), "File does not exist: {file_path:?}");

        assert_eq!(image_data_back.cols(), 258);
        assert_eq!(image_data_back.rows(), 195);
        assert_eq!(image_data_back.num_channels(), 3);

        Ok(())
    }

    #[test]
    fn read_write_png_rgb16() -> Result<(), IoError> {
        let tmp_dir = tempfile::tempdir()?;
        create_dir_all(tmp_dir.path())?;

        let file_path = tmp_dir.path().join("rgb16.png");
        let image_data = read_image_png_rgb16("../../tests/data/rgb16.png")?;
        write_image_png_rgb16(&file_path, &image_data)?;

        let image_data_back = read_image_png_rgb16(&file_path)?;
        assert!(file_path.exists(), "File does not exist: {file_path:?}");

        assert_eq!(image_data_back.cols(), 32);
        assert_eq!(image_data_back.rows(), 32);
        assert_eq!(image_data_back.num_channels(), 3);

        Ok(())
    }

    #[test]
    fn decode_png() -> Result<(), IoError> {
        let bytes = read("../../tests/data/dog-rgb8.png")?;
        let mut image = Rgb8::from_size_val([258, 195].into(), 0)?;
        decode_image_png_rgb8(&bytes, &mut image)?;

        assert_eq!(image.cols(), 258);
        assert_eq!(image.rows(), 195);
        assert_eq!(image.num_channels(), 3);

        Ok(())
    }

    // -----------------------------------------------------------------
    // In-memory encode round-trips: encode → decode equals the source.
    // -----------------------------------------------------------------

    #[test]
    fn encode_decode_png_rgb8_roundtrip() -> Result<(), IoError> {
        let src = read_image_png_rgb8("../../tests/data/dog-rgb8.png")?;
        let mut buffer = Vec::new();
        encode_image_png_rgb8(&src, &mut buffer, None)?;
        assert!(!buffer.is_empty());
        // PNG magic header
        assert_eq!(&buffer[..8], b"\x89PNG\r\n\x1a\n");

        let mut decoded = Rgb8::from_size_val(src.size(), 0)?;
        decode_image_png_rgb8(&buffer, &mut decoded)?;
        assert_eq!(decoded.size(), src.size());
        assert_eq!(decoded.as_slice(), src.as_slice());
        Ok(())
    }

    #[test]
    fn encode_decode_png_rgba8_roundtrip() -> Result<(), IoError> {
        // Synthesize an RGBA8 image (no fixture in tests/data for this layout).
        let mut data = vec![0u8; 16 * 16 * 4];
        for (i, b) in data.iter_mut().enumerate() {
            *b = (i % 251) as u8;
        }
        let src = Rgba8::from_size_vec([16, 16].into(), data)?;

        let mut buffer = Vec::new();
        encode_image_png_rgba8(&src, &mut buffer, None)?;
        assert_eq!(&buffer[..8], b"\x89PNG\r\n\x1a\n");

        let mut decoded = Rgba8::from_size_val(src.size(), 0)?;
        decode_image_png_rgba8(&buffer, &mut decoded)?;
        assert_eq!(decoded.as_slice(), src.as_slice());
        Ok(())
    }

    #[test]
    fn encode_decode_png_gray8_roundtrip() -> Result<(), IoError> {
        let src = read_image_png_mono8("../../tests/data/dog.png")?;
        let mut buffer = Vec::new();
        encode_image_png_gray8(&src, &mut buffer, None)?;
        assert_eq!(&buffer[..8], b"\x89PNG\r\n\x1a\n");

        let mut decoded = Gray8::from_size_val(src.size(), 0)?;
        decode_image_png_mono8(&buffer, &mut decoded)?;
        assert_eq!(decoded.as_slice(), src.as_slice());
        Ok(())
    }

    #[test]
    fn encode_decode_png_rgb16_roundtrip() -> Result<(), IoError> {
        let src = read_image_png_rgb16("../../tests/data/rgb16.png")?;
        let mut buffer = Vec::new();
        encode_image_png_rgb16(&src, &mut buffer, None)?;
        assert_eq!(&buffer[..8], b"\x89PNG\r\n\x1a\n");

        let mut decoded = Rgb16::from_size_val(src.size(), 0)?;
        decode_image_png_rgb16(&buffer, &mut decoded)?;
        assert_eq!(decoded.as_slice(), src.as_slice());
        Ok(())
    }

    #[test]
    fn encode_decode_png_gray16_roundtrip() -> Result<(), IoError> {
        // Depth-style synthetic data: smooth gradient + a sharp object, in mm.
        let (w, h) = (64usize, 48usize);
        let mut data = vec![0u16; w * h];
        for y in 0..h {
            for x in 0..w {
                data[y * w + x] = 1000 + (x as u16) * 8 + (y as u16) * 4;
            }
        }
        // Object discontinuity — the kind of edge that JPEG would smear.
        for y in 10..20 {
            for x in 20..40 {
                data[y * w + x] = 500;
            }
        }
        let src = Gray16::from_size_vec([w, h].into(), data)?;

        let mut buffer = Vec::new();
        encode_image_png_gray16(&src, &mut buffer, None)?;
        assert_eq!(&buffer[..8], b"\x89PNG\r\n\x1a\n");
        // Lossless u16 round-trip is the whole point of this codec for depth.
        let mut decoded = Gray16::from_size_val(src.size(), 0)?;
        decode_image_png_mono16(&buffer, &mut decoded)?;
        assert_eq!(decoded.as_slice(), src.as_slice());
        Ok(())
    }

    #[test]
    fn encode_png_buffer_reuse() -> Result<(), IoError> {
        // A streaming caller (e.g. a recorder) reuses one buffer across many
        // encodes — verify two back-to-back encodes into the same Vec produce
        // independently-decodable PNGs after `clear()`.
        let src = read_image_png_rgb8("../../tests/data/dog-rgb8.png")?;
        let mut buffer = Vec::with_capacity(64 * 1024);

        encode_image_png_rgb8(&src, &mut buffer, None)?;
        let cap_after_first = buffer.capacity();

        buffer.clear();
        encode_image_png_rgb8(&src, &mut buffer, None)?;

        // Capacity should not have grown on the second encode (allocation reuse).
        assert!(buffer.capacity() <= cap_after_first.max(buffer.len()));
        let mut decoded = Rgb8::from_size_val(src.size(), 0)?;
        decode_image_png_rgb8(&buffer, &mut decoded)?;
        assert_eq!(decoded.as_slice(), src.as_slice());
        Ok(())
    }
}