kornia-io 0.2.0

use crate::{
    error::IoError,
    limits::{alloc_image, check_image_dimensions},
};
use image_webp::{ColorType, WebPDecoder, WebPEncoder};
use kornia_image::{
    color_spaces::{Gray8, Rgb8, Rgba8},
    Image, ImageLayout, ImageSize, PixelFormat,
};

use std::{
    fs,
    io::{BufReader, Cursor},
    path::Path,
};

// BT.601 luma weights, fixed-point /256: R=0.299, G=0.587, B=0.114.
const GRAY_R: u32 = 77;
const GRAY_G: u32 = 150;
const GRAY_B: u32 = 29;

/// Read a WEBP image as grayscale (Gray8).
///
/// WebP has no native grayscale encoding; the file is decoded as RGB(A) and
/// then converted to luma using BT.601 weights.
///
/// # Arguments
///
/// * `file_path` - The path to the WEBP file.
///
/// # Returns
///
/// A grayscale image (Gray8).
pub fn read_image_webp_gray8(file_path: impl AsRef<Path>) -> Result<Gray8, IoError> {
    let mut decoder = open_webp(file_path)?;
    let (width, height) = decoder.dimensions();
    let mut gray = alloc_image(ImageSize {
        width: width as usize,
        height: height as usize,
    })?;
    decode_webp_gray_into(&mut decoder, &mut gray)?;
    Ok(Gray8(gray))
}

/// Read a WEBP image as RGB8.
///
/// Returns an error if the file contains an alpha channel.
///
/// # Arguments
///
/// * `file_path` - The path to the WEBP file.
///
/// # Returns
///
/// A RGB8 typed image.
pub fn read_image_webp_rgb8(file_path: impl AsRef<Path>) -> Result<Rgb8, IoError> {
    let mut decoder = open_webp(file_path)?;
    if decoder.has_alpha() {
        return Err(IoError::WebpDecodingError(
            image_webp::DecodingError::InvalidParameter(
                "file has alpha channel; use read_image_webp_rgba8".to_string(),
            ),
        ));
    }
    Ok(Rgb8(decode_webp_new(&mut decoder)?))
}

/// Read a WEBP image as RGBA8.
///
/// Returns an error if the file has no alpha channel.
///
/// # Arguments
///
/// * `file_path` - The path to the WEBP file.
///
/// # Returns
///
/// A RGBA8 typed image.
pub fn read_image_webp_rgba8(file_path: impl AsRef<Path>) -> Result<Rgba8, IoError> {
    let mut decoder = open_webp(file_path)?;
    if !decoder.has_alpha() {
        return Err(IoError::WebpDecodingError(
            image_webp::DecodingError::InvalidParameter(
                "file has no alpha channel; use read_image_webp_rgb8".to_string(),
            ),
        ));
    }
    Ok(Rgba8(decode_webp_new(&mut decoder)?))
}

/// Decodes a WEBP image as RGB8 from raw bytes.
///
/// Errors if `src` contains an alpha channel or if `dst` dimensions do not match.
///
/// # Arguments
///
/// - `src` - Raw bytes of the webp file
/// - `dst` - A mutable reference to your `Rgb8` image
pub fn decode_image_webp_rgb8(src: &[u8], dst: &mut Image<u8, 3>) -> Result<(), IoError> {
    decode_webp_impl::<3>(src, dst, false)
}

/// Decodes a WEBP image as RGBA8 from raw bytes.
///
/// Errors if `src` has no alpha channel or if `dst` dimensions do not match.
///
/// # Arguments
///
/// - `src` - Raw bytes of the webp file
/// - `dst` - A mutable reference to your `Rgba8` image
pub fn decode_image_webp_rgba8(src: &[u8], dst: &mut Image<u8, 4>) -> Result<(), IoError> {
    decode_webp_impl::<4>(src, dst, true)
}

/// Decodes a WEBP image as Gray8 from raw bytes.
///
/// WebP has no native grayscale encoding; the file is decoded as RGB(A) and
/// converted to luma using BT.601 weights.
///
/// # Arguments
///
/// - `src` - Raw bytes of the webp file
/// - `dst` - A mutable reference to your `Gray8` image
pub fn decode_image_webp_gray8(src: &[u8], dst: &mut Image<u8, 1>) -> Result<(), IoError> {
    let mut decoder = WebPDecoder::new(Cursor::new(src))?;
    let (width, height) = decoder.dimensions();
    if [width as usize, height as usize] != [dst.width(), dst.height()] {
        return Err(IoError::DecodeMismatchResolution(
            height as usize,
            width as usize,
            dst.height(),
            dst.width(),
        ));
    }

    decode_webp_gray_into(&mut decoder, dst)
}

/// Decodes WEBP image metadata from raw bytes without decoding pixel data.
///
/// # Arguments
///
/// - `src` - Raw bytes of the WEBP file
///
/// # Returns
///
/// An `ImageLayout` containing the image metadata (size, channels, pixel format).
/// Channel count is 3 (RGB) or 4 (RGBA); WebP has no native grayscale encoding.
pub fn decode_image_webp_layout(src: &[u8]) -> Result<ImageLayout, IoError> {
    let decoder = WebPDecoder::new(Cursor::new(src))?;
    let (width, height) = decoder.dimensions();
    check_image_dimensions(width as usize, height as usize)?;
    let channels: u8 = if decoder.has_alpha() { 4 } else { 3 };
    Ok(ImageLayout::new(
        ImageSize {
            width: width as usize,
            height: height as usize,
        },
        channels,
        PixelFormat::U8,
    ))
}

// Decodes a WEBP image into a pre-allocated Image buffer, validating channel count and size.
fn decode_webp_impl<const C: usize>(
    src: &[u8],
    dst: &mut Image<u8, C>,
    expect_alpha: bool,
) -> Result<(), IoError> {
    let mut decoder = WebPDecoder::new(Cursor::new(src))?;

    let (width, height) = decoder.dimensions();
    if [width as usize, height as usize] != [dst.width(), dst.height()] {
        return Err(IoError::DecodeMismatchResolution(
            height as usize,
            width as usize,
            dst.height(),
            dst.width(),
        ));
    }

    if decoder.has_alpha() != expect_alpha {
        return Err(IoError::WebpDecodingError(
            image_webp::DecodingError::InvalidParameter(format!(
                "channel mismatch: file has_alpha={} but dst expects {} channels",
                decoder.has_alpha(),
                C
            )),
        ));
    }

    let expected_len = (width as usize) * (height as usize) * C;
    let dst_slice = dst.as_slice_mut();
    if dst_slice.len() != expected_len {
        return Err(IoError::InvalidBufferSize(dst_slice.len(), expected_len));
    }

    decoder.read_image(dst_slice)?;
    Ok(())
}

// Opens a WebP file after validating that it exists and has a `.webp` extension.
fn open_webp(file_path: impl AsRef<Path>) -> Result<WebPDecoder<BufReader<fs::File>>, IoError> {
    let file_path = file_path.as_ref();
    if !file_path.exists() {
        return Err(IoError::FileDoesNotExist(file_path.to_path_buf()));
    }
    match file_path.extension() {
        Some(ext) if ext == "webp" => {}
        _ => return Err(IoError::InvalidFileExtension(file_path.to_path_buf())),
    }

    let file = fs::File::open(file_path)?;
    Ok(WebPDecoder::new(BufReader::new(file))?)
}

// Decodes into a newly allocated `C`-channel image. `C` must be 4 if the file has an
// alpha channel and 3 otherwise; the decoder rejects a mismatched buffer length.
fn decode_webp_new<R: std::io::BufRead + std::io::Seek, const C: usize>(
    decoder: &mut WebPDecoder<R>,
) -> Result<Image<u8, C>, IoError> {
    let (width, height) = decoder.dimensions();
    let mut img = alloc_image::<u8, C>(ImageSize {
        width: width as usize,
        height: height as usize,
    })?;
    decoder.read_image(img.as_slice_mut())?;
    Ok(img)
}

#[inline]
fn luma_from_rgb(r: u8, g: u8, b: u8) -> u8 {
    ((r as u32 * GRAY_R + g as u32 * GRAY_G + b as u32 * GRAY_B) >> 8) as u8
}

// Decodes as RGB(A) and converts to luma. `dst` must match the decoded size, so
// each RGB(A) pixel maps to exactly one luma sample.
fn decode_webp_gray_into<R: std::io::BufRead + std::io::Seek>(
    decoder: &mut WebPDecoder<R>,
    dst: &mut Image<u8, 1>,
) -> Result<(), IoError> {
    fn convert<const C: usize>(src: &Image<u8, C>, dst: &mut Image<u8, 1>) {
        for (d, c) in dst
            .as_slice_mut()
            .iter_mut()
            .zip(src.as_slice().chunks_exact(C))
        {
            *d = luma_from_rgb(c[0], c[1], c[2]);
        }
    }
    if decoder.has_alpha() {
        convert(&decode_webp_new::<_, 4>(decoder)?, dst);
    } else {
        convert(&decode_webp_new::<_, 3>(decoder)?, dst);
    }
    Ok(())
}

/// Encodes the given RGB8 image to WEBP bytes (VP8L lossless).
///
/// # Arguments
///
/// - `image` - The RGB image to encode
/// - `buffer` - A mutable buffer to write the WEBP bytes into. Existing contents are preserved;
///   the encoded data is appended.
pub fn encode_image_webp_rgb8(image: &Image<u8, 3>, buffer: &mut Vec<u8>) -> Result<(), IoError> {
    WebPEncoder::new(buffer).encode(
        image.as_slice(),
        image.width() as u32,
        image.height() as u32,
        ColorType::Rgb8,
    )?;
    Ok(())
}

/// Encodes the given RGBA8 image to WEBP bytes (VP8L lossless).
///
/// # Arguments
///
/// - `image` - The RGBA8 image to encode
/// - `buffer` - A mutable buffer to write the WEBP bytes into. Existing contents are preserved;
///   the encoded data is appended.
pub fn encode_image_webp_rgba8(image: &Image<u8, 4>, buffer: &mut Vec<u8>) -> Result<(), IoError> {
    WebPEncoder::new(buffer).encode(
        image.as_slice(),
        image.width() as u32,
        image.height() as u32,
        ColorType::Rgba8,
    )?;
    Ok(())
}

/// Encodes the given Gray8 image to WEBP bytes (VP8L lossless).
///
/// # Arguments
///
/// - `image` - The Gray8 image to encode
/// - `buffer` - A mutable buffer to write the WEBP bytes into. Existing contents are preserved;
///   the encoded data is appended.
pub fn encode_image_webp_gray8(image: &Image<u8, 1>, buffer: &mut Vec<u8>) -> Result<(), IoError> {
    WebPEncoder::new(buffer).encode(
        image.as_slice(),
        image.width() as u32,
        image.height() as u32,
        ColorType::L8,
    )?;
    Ok(())
}

/// Writes the given Gray8 image to the given file path as WEBP.
///
/// # Arguments
///
/// - `file_path` - The path to the WEBP image.
/// - `image` - The grayscale image to write
pub fn write_image_webp_gray8(
    file_path: impl AsRef<Path>,
    image: &Image<u8, 1>,
) -> Result<(), IoError> {
    write_image_webp_impl(file_path, image, ColorType::L8)
}

/// Writes the given RGB8 image to the given file path as WEBP.
///
/// # Arguments
///
/// - `file_path` - The path to the WEBP image.
/// - `image` - The rgb8 image to write
pub fn write_image_webp_rgb8(
    file_path: impl AsRef<Path>,
    image: &Image<u8, 3>,
) -> Result<(), IoError> {
    write_image_webp_impl(file_path, image, ColorType::Rgb8)
}

/// Writes the given RGBA8 image to the given file path as WEBP.
///
/// # Arguments
///
/// - `file_path` - The path to the WEBP image.
/// - `image` - The rgba8 image to write
pub fn write_image_webp_rgba8(
    file_path: impl AsRef<Path>,
    image: &Image<u8, 4>,
) -> Result<(), IoError> {
    write_image_webp_impl(file_path, image, ColorType::Rgba8)
}

fn write_image_webp_impl<const N: usize>(
    file_path: impl AsRef<Path>,
    image: &Image<u8, N>,
    color_type: ColorType,
) -> Result<(), IoError> {
    let file = fs::File::create(file_path)?;
    let writer = std::io::BufWriter::new(file);
    WebPEncoder::new(writer).encode(
        image.as_slice(),
        image.width() as u32,
        image.height() as u32,
        color_type,
    )?;
    Ok(())
}

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

    // A 44-byte WebP whose VP8X header declares a 65535x65535 canvas.
    fn webp_bomb() -> Vec<u8> {
        let mut vp8x = [0u8; 10];
        vp8x[4..7].copy_from_slice(&65534u32.to_le_bytes()[..3]);
        vp8x[7..10].copy_from_slice(&65534u32.to_le_bytes()[..3]);
        let mut body = Vec::new();
        body.extend_from_slice(b"VP8X");
        body.extend_from_slice(&10u32.to_le_bytes());
        body.extend_from_slice(&vp8x);
        let vp8l = [0x2fu8, 0, 0, 0, 0];
        body.extend_from_slice(b"VP8L");
        body.extend_from_slice(&(vp8l.len() as u32).to_le_bytes());
        body.extend_from_slice(&vp8l);
        body.push(0);
        let mut out = b"RIFF".to_vec();
        out.extend_from_slice(&((body.len() + 4) as u32).to_le_bytes());
        out.extend_from_slice(b"WEBP");
        out.extend_from_slice(&body);
        out
    }

    #[test]
    fn rejects_decompression_bomb() -> Result<(), Box<dyn std::error::Error>> {
        let bomb = webp_bomb();
        assert!(matches!(
            decode_image_webp_layout(&bomb),
            Err(IoError::ImageTooLarge { .. })
        ));
        let dir = tempfile::tempdir()?;
        let path = dir.path().join("bomb.webp");
        std::fs::write(&path, &bomb)?;
        assert!(matches!(
            read_image_webp_rgb8(&path),
            Err(IoError::ImageTooLarge { .. })
        ));
        Ok(())
    }

    #[test]
    fn test_read_webp_rgb8() -> Result<(), IoError> {
        let image = read_image_webp_rgb8("../../tests/data/fire.webp")?;
        assert_eq!(image.cols(), 320);
        assert_eq!(image.rows(), 235);
        Ok(())
    }

    #[test]
    fn test_read_webp_gray8() -> Result<(), IoError> {
        let image = read_image_webp_gray8("../../tests/data/fire.webp")?;
        assert_eq!(image.cols(), 320);
        assert_eq!(image.rows(), 235);
        Ok(())
    }

    #[test]
    fn test_decode_webp() -> Result<(), IoError> {
        let bytes = read("../../tests/data/fire.webp")?;
        let mut image = Rgb8::from_size_val([320, 235].into(), 0)?;
        decode_image_webp_rgb8(&bytes, &mut image)?;

        assert_eq!(image.cols(), 320);
        assert_eq!(image.rows(), 235);
        assert_eq!(image.num_channels(), 3);

        Ok(())
    }

    #[test]
    fn test_decode_webp_layout_size() -> Result<(), IoError> {
        let bytes = read("../../tests/data/fire.webp")?;
        let layout = decode_image_webp_layout(bytes.as_slice())?;
        assert_eq!(layout.image_size.width, 320);
        assert_eq!(layout.image_size.height, 235);
        assert_eq!(layout.channels, 3);
        Ok(())
    }

    #[test]
    fn read_write_webp_rgb8() -> Result<(), IoError> {
        let tmp_dir = tempfile::tempdir()?;
        let file_path = tmp_dir.path().join("fire_write_rgb8.webp");
        let image_data = read_image_webp_rgb8("../../tests/data/fire.webp")?;
        write_image_webp_rgb8(&file_path, &image_data)?;

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

        assert_eq!(image_data_back.cols(), 320);
        assert_eq!(image_data_back.rows(), 235);
        assert_eq!(image_data_back.num_channels(), 3);
        assert_eq!(image_data.as_slice(), image_data_back.as_slice());

        Ok(())
    }

    #[test]
    fn read_write_webp_rgba8() -> Result<(), IoError> {
        let tmp_dir = tempfile::tempdir()?;
        let file_path = tmp_dir.path().join("synthetic_rgba8.webp");

        let w = 16;
        let h = 8;
        let mut pixels = Vec::with_capacity(w * h * 4);
        for y in 0..h {
            for x in 0..w {
                pixels.extend_from_slice(&[x as u8, y as u8, (x + y) as u8, 0x80]);
            }
        }
        let src = Rgba8::from_size_vec([w, h].into(), pixels)?;
        write_image_webp_rgba8(&file_path, &src)?;

        let decoded = read_image_webp_rgba8(&file_path)?;
        assert_eq!(decoded.cols(), w);
        assert_eq!(decoded.rows(), h);
        assert_eq!(decoded.num_channels(), 4);
        assert_eq!(decoded.as_slice(), src.as_slice());

        Ok(())
    }

    #[test]
    fn rejects_non_webp_extension() {
        match read_image_webp_rgb8("../../tests/data/dog.jpeg") {
            Err(IoError::InvalidFileExtension(_)) => {}
            other => panic!("expected InvalidFileExtension, got {:?}", other.err()),
        }
    }

    #[test]
    fn rgb_reader_rejects_rgba_file() -> Result<(), IoError> {
        // Encode a synthetic RGBA webp, then try to read it with the RGB reader.
        let tmp_dir = tempfile::tempdir()?;
        let file_path = tmp_dir.path().join("rgba_only.webp");
        let w = 4;
        let h = 4;
        let pixels = vec![0xAAu8; w * h * 4];
        let src = Rgba8::from_size_vec([w, h].into(), pixels)?;
        write_image_webp_rgba8(&file_path, &src)?;

        match read_image_webp_rgb8(&file_path) {
            Err(IoError::WebpDecodingError(_)) => Ok(()),
            Err(other) => panic!("expected WebpDecodingError, got {:?}", other),
            Ok(_) => panic!("expected error, got Ok"),
        }
    }
}