img-optimize 0.1.0

use std::io::Cursor;
use cfg_if::cfg_if;
use crate::structs::*;

pub fn encode_auto_select<'a>(im_info: &ImageInfo, config: &CompressionConfig) -> &'a str {
    #[allow(unused_mut)] let mut support_webp = true;
    #[allow(unused_mut)] let mut support_png = true;
    #[allow(unused_mut)] let mut support_jpeg_xl = true;

    cfg_if! {
        if #[cfg(not(feature="webp"))] {
            support_webp = false;
        }
    }
    cfg_if! {
        if #[cfg(not(feature="png"))] {
            support_png = false;
        }
    }
    cfg_if! {
        if #[cfg(not(feature="jpegxl"))] {
            support_jpeg_xl = false;
        }
    }

    match im_info.pixels {
        ImagePixels::U8(_) => {
            if support_webp && is_webp_supported_image(&im_info, &config) {
                // If 8-bit color and RGB or RGBA, use webp lossless
                "webp"
            }
            else if support_png && is_png_supported_image(&im_info, &config) {
                "png"
            }
            else
            {
                ""
            }
        },
        ImagePixels::U16(_) => {
            if support_png && is_png_supported_image(&im_info, &config) {
                "png"
            }
            else
            {
                ""
            }
        },
        _ => {
            ""
        }
    }
}

// Returns: (file_extension, file_content)
pub fn encode_auto(im_info: ImageInfo, config: CompressionConfig) -> (String, Vec<u8>) {
    let decide = encode_auto_select(&im_info, &config);

    match decide {
        #[cfg(feature="webp")]
        "webp" => {
            let webp_memory = encode_lossless_webp(im_info);
            ("webp".to_owned(), webp_memory)
        },
        #[cfg(feature="png")]
        "png" => {
            let png_memory = encode_png(im_info, config);
            ("png".to_owned(), png_memory)
        },
        _ => {
            panic!("Unsupported image format");
        }
    }
}

fn is_webp_supported_image(im_info: &ImageInfo, config: &CompressionConfig) -> bool {
    config.allow_webp &&
        (im_info.bit_depth == 8) &&
        (im_info.color_type == 2 || im_info.color_type == 6)
}

fn is_png_supported_image(im_info: &ImageInfo, config: &CompressionConfig) -> bool {
    config.allow_png &&
        (im_info.bit_depth == 1 || im_info.bit_depth == 2 || im_info.bit_depth == 4 || im_info.bit_depth == 8 || im_info.bit_depth == 16) &&
        (im_info.color_type == 0 || im_info.color_type == 2 || /* im_info.color_type == 3 || */ im_info.color_type == 4 || im_info.color_type == 6)
}

#[cfg(feature="webp")]
pub fn encode_lossless_webp(im_info: ImageInfo) -> Vec<u8> {
    let bit_per_pixel = match im_info.color_type {
        2 => 3,
        6 => 4,
        _ => panic!("Unsupported color type"),
    };

    let pixel_layout = match im_info.color_type {
        2 => webp::PixelLayout::Rgb,
        6 => webp::PixelLayout::Rgba,
        _ => panic!("Unsupported color type"),
    };

    let mut image_pixels = vec![0; im_info.width as usize * im_info.height as usize * bit_per_pixel];

    match im_info.pixels {
        ImagePixels::U8(pixels) => {
            for i in 0..pixels.len() {
                image_pixels[i] = pixels[i];
            }
        },
        _ => panic!("Unsupported pixel type"),
    }

    let encoder = webp::Encoder::new(&*image_pixels, pixel_layout, im_info.width, im_info.height);

    // ToDo: Add lossy support
    let webp_memory = encoder.encode_lossless();

    webp_memory.to_vec()
}

#[cfg(feature="png")]
pub fn encode_png(im_info: ImageInfo, compression_config: CompressionConfig) -> Vec<u8> {
    cfg_if! {
        if #[cfg(feature="oxipng")] {
            return encode_png_oxipng(im_info, compression_config);
        } else {
            return encode_png_image_rs(im_info);
        }
    }
}

#[cfg(feature="png")]
pub fn encode_png_image_rs(im_info: ImageInfo) -> Vec<u8> {
    let mut output_vec: Vec<u8> = Vec::new();
    {
        let output_cursor = Cursor::new(&mut output_vec);

        let mut encoder = png::Encoder::new(output_cursor, im_info.width, im_info.height);

        let color_type = match im_info.color_type {
            0 => png::ColorType::Grayscale,
            2 => png::ColorType::Rgb,
            3 => png::ColorType::Indexed,
            4 => png::ColorType::GrayscaleAlpha,
            6 => png::ColorType::Rgba,
            _ => panic!("Unsupported color type"),
        };

        encoder.set_color(color_type);

        let color_depth = match im_info.bit_depth {
            1 => png::BitDepth::One,
            2 => png::BitDepth::Two,
            4 => png::BitDepth::Four,
            8 => png::BitDepth::Eight,
            16 => png::BitDepth::Sixteen,
            _ => panic!("Unsupported bit_depth"),
        };

        encoder.set_depth(color_depth);

        let mut pixels;

        match im_info.pixels {
            ImagePixels::U8(pxl) => {
                pixels = pxl;
            },
            ImagePixels::U16(pxl) => {
                pixels = Vec::new();
                for i in 0..pxl.len() {
                    pixels.push((pxl[i] >> 8) as u8);
                    pixels.push((pxl[i] & 0xFF) as u8);
                }
            },
            _ => panic!("Unsupported pixel type"),
        }

        let mut writer = encoder.write_header().unwrap();

        // ToDo: This might be not support 16-bit image.
        writer.write_image_data(&*pixels).unwrap();
    }

    let mut output: Vec<u8> = vec![0; output_vec.len()];
    for i in 0..output.len() {
        output[i] = (&*output_vec)[i];
    }
    output
}

#[cfg(feature="oxipng")]
pub fn encode_png_oxipng(im_info: ImageInfo, compression_config: CompressionConfig) -> Vec<u8> {
    let color_type = match im_info.color_type {
        0 => oxipng::ColorType::Grayscale {
            transparent_shade: None,
        },
        2 => oxipng::ColorType::RGB {
            transparent_color: None,
        },
        4 => oxipng::ColorType::GrayscaleAlpha,
        6 => oxipng::ColorType::RGBA,
        _ => panic!("Unsupported color type"),
    };

    let bit_depth = match im_info.bit_depth {
        1 => oxipng::BitDepth::One,
        2 => oxipng::BitDepth::Two,
        4 => oxipng::BitDepth::Four,
        8 => oxipng::BitDepth::Eight,
        16 => oxipng::BitDepth::Sixteen,
        _ => panic!("Unsupported bit depth"),
    };

    let pixels;
    match im_info.pixels {
        ImagePixels::U8(pixels8) => {
            pixels = pixels8;
        },
        ImagePixels::U16(pixels16) => {
            let mut pixels8 = vec![0; pixels16.len() * 2];
            for i in 0..pixels16.len() {
                pixels8[i * 2] = (pixels16[i] >> 8) as u8;
                pixels8[i * 2 + 1] = pixels16[i] as u8;
            }
            pixels = pixels8;
        },
        _ => panic!("Unsupported pixel type"),
    }

    let oxipng_img = oxipng::RawImage::new(
        im_info.width,
        im_info.height,
        color_type,
        bit_depth,
        pixels,
    ).unwrap();

    let mut options = oxipng::Options::max_compression();
    if compression_config.allow_oxipng_zopfli {
        options.deflate = oxipng::Deflaters::Zopfli {
            iterations: std::num::NonZeroU8::new(15).unwrap(), // This is default of oxipng cli.
        };
    }

    let optimized = oxipng_img.create_optimized_png(&options).unwrap();
    optimized
}