pixstrip/pixstrip-core/src/encoder.rs

use std::io::Cursor;
use std::path::Path;

use crate::error::{PixstripError, Result};
use crate::types::{ImageFormat, QualityPreset};

#[derive(Debug, Clone, Default)]
pub struct EncoderOptions {
    pub progressive_jpeg: bool,
    pub avif_speed: u8,
    /// Output DPI (0 means don't set / use default)
    pub output_dpi: u32,
}

pub struct OutputEncoder {
    pub options: EncoderOptions,
}

impl OutputEncoder {
    pub fn new() -> Self {
        Self {
            options: EncoderOptions {
                progressive_jpeg: false,
                avif_speed: 6,
                output_dpi: 0,
            },
        }
    }

    pub fn with_options(options: EncoderOptions) -> Self {
        Self { options }
    }

    pub fn encode(
        &self,
        img: &image::DynamicImage,
        format: ImageFormat,
        quality: Option<u8>,
    ) -> Result<Vec<u8>> {
        match format {
            ImageFormat::Jpeg => self.encode_jpeg(img, quality.unwrap_or(85)),
            ImageFormat::Png => self.encode_png(img, quality.unwrap_or(3)),
            ImageFormat::WebP => self.encode_webp(img, quality.unwrap_or(80)),
            ImageFormat::Avif => self.encode_avif(img, quality.unwrap_or(80)),
            ImageFormat::Gif => self.encode_fallback(img, image::ImageFormat::Gif),
            ImageFormat::Tiff => self.encode_fallback(img, image::ImageFormat::Tiff),
            ImageFormat::Bmp => self.encode_fallback(img, image::ImageFormat::Bmp),
        }
    }

    pub fn encode_to_file(
        &self,
        img: &image::DynamicImage,
        path: &Path,
        format: ImageFormat,
        quality: Option<u8>,
    ) -> Result<()> {
        let bytes = self.encode(img, format, quality)?;
        if let Some(parent) = path.parent() {
            std::fs::create_dir_all(parent).map_err(PixstripError::Io)?;
        }
        std::fs::write(path, &bytes).map_err(PixstripError::Io)?;
        Ok(())
    }

    pub fn quality_for_format(&self, format: ImageFormat, preset: &QualityPreset) -> u8 {
        match format {
            ImageFormat::Jpeg => preset.jpeg_quality(),
            ImageFormat::WebP => preset.webp_quality() as u8,
            ImageFormat::Avif => preset.avif_quality() as u8,
            ImageFormat::Png => preset.png_level(),
            ImageFormat::Gif | ImageFormat::Tiff | ImageFormat::Bmp => preset.jpeg_quality(),
        }
    }

    fn encode_jpeg(&self, img: &image::DynamicImage, quality: u8) -> Result<Vec<u8>> {
        let rgb = img.to_rgb8();
        let (width, height) = (rgb.width() as usize, rgb.height() as usize);
        let pixels = rgb.as_raw();

        let mut comp = mozjpeg::Compress::new(mozjpeg::ColorSpace::JCS_RGB);
        comp.set_size(width, height);
        comp.set_quality(quality as f32);
        if self.options.progressive_jpeg {
            comp.set_progressive_mode();
        }
        if self.options.output_dpi > 0 {
            comp.set_pixel_density(mozjpeg::PixelDensity {
                unit: mozjpeg::PixelDensityUnit::Inches,
                x: self.options.output_dpi.min(65535) as u16,
                y: self.options.output_dpi.min(65535) as u16,
            });
        }

        let mut output = Vec::new();
        let mut started = comp.start_compress(&mut output).map_err(|e| PixstripError::Processing {
            operation: "jpeg_encode".into(),
            reason: e.to_string(),
        })?;

        let row_stride = width * 3;
        for y in 0..height {
            let start = y * row_stride;
            let end = start + row_stride;
            started.write_scanlines(&pixels[start..end]).map_err(|e| PixstripError::Processing {
                operation: "jpeg_scanline".into(),
                reason: e.to_string(),
            })?;
        }

        started.finish().map_err(|e| PixstripError::Processing {
            operation: "jpeg_encode".into(),
            reason: e.to_string(),
        })?;

        Ok(output)
    }

    fn encode_png(&self, img: &image::DynamicImage, level: u8) -> Result<Vec<u8>> {
        let mut buf = Vec::new();
        let cursor = Cursor::new(&mut buf);
        let rgba = img.to_rgba8();
        let encoder = image::codecs::png::PngEncoder::new(cursor);
        image::ImageEncoder::write_image(
            encoder,
            rgba.as_raw(),
            rgba.width(),
            rgba.height(),
            image::ExtendedColorType::Rgba8,
        )
        .map_err(|e| PixstripError::Processing {
            operation: "png_encode".into(),
            reason: e.to_string(),
        })?;

        // Insert pHYs chunk for DPI if requested (before oxipng, which preserves it)
        if self.options.output_dpi > 0 {
            buf = insert_png_phys_chunk(&buf, self.options.output_dpi);
        }

        let mut opts = oxipng::Options::default();
        opts.optimize_alpha = true;
        opts.deflater = oxipng::Deflater::Libdeflater { compression: level.clamp(1, 12) };

        let optimized = oxipng::optimize_from_memory(&buf, &opts)
            .map_err(|e| PixstripError::Processing {
                operation: "png_optimize".into(),
                reason: e.to_string(),
            })?;

        Ok(optimized)
    }

    fn encode_webp(&self, img: &image::DynamicImage, quality: u8) -> Result<Vec<u8>> {
        let encoder = webp::Encoder::from_image(img).map_err(|e| PixstripError::Processing {
            operation: "webp_encode".into(),
            reason: e.to_string(),
        })?;
        let mem = encoder.encode(quality as f32);
        Ok(mem.to_vec())
    }

    fn encode_avif(&self, img: &image::DynamicImage, quality: u8) -> Result<Vec<u8>> {
        let mut buf = Vec::new();
        let cursor = Cursor::new(&mut buf);
        let rgba = img.to_rgba8();
        let speed = self.options.avif_speed.clamp(0, 10);
        let encoder = image::codecs::avif::AvifEncoder::new_with_speed_quality(
            cursor,
            speed,
            quality,
        );
        image::ImageEncoder::write_image(
            encoder,
            rgba.as_raw(),
            rgba.width(),
            rgba.height(),
            image::ExtendedColorType::Rgba8,
        )
        .map_err(|e| PixstripError::Processing {
            operation: "avif_encode".into(),
            reason: e.to_string(),
        })?;
        Ok(buf)
    }

    fn encode_fallback(
        &self,
        img: &image::DynamicImage,
        format: image::ImageFormat,
    ) -> Result<Vec<u8>> {
        let mut buf = Vec::new();
        let cursor = Cursor::new(&mut buf);
        img.write_to(cursor, format).map_err(|e| PixstripError::Processing {
            operation: "encode".into(),
            reason: e.to_string(),
        })?;
        Ok(buf)
    }
}

/// Insert a pHYs chunk into PNG data to set DPI.
/// The pHYs chunk must appear before the first IDAT chunk.
/// DPI is converted to pixels per meter (1 inch = 0.0254 meters).
fn insert_png_phys_chunk(png_data: &[u8], dpi: u32) -> Vec<u8> {
    // Not a valid PNG (too short for signature) - return as-is
    if png_data.len() < 8 {
        return png_data.to_vec();
    }

    // PNG pixels per meter = DPI / 0.0254
    let ppm = (dpi as f64 / 0.0254).round() as u32;

    // Build the pHYs chunk:
    // 4 bytes: data length (9)
    // 4 bytes: chunk type "pHYs"
    // 4 bytes: pixels per unit X
    // 4 bytes: pixels per unit Y
    // 1 byte: unit (1 = meter)
    // 4 bytes: CRC32 of type + data
    let mut chunk_data = Vec::with_capacity(9);
    chunk_data.extend_from_slice(&ppm.to_be_bytes()); // X pixels per unit
    chunk_data.extend_from_slice(&ppm.to_be_bytes()); // Y pixels per unit
    chunk_data.push(1); // unit = meter

    let mut crc_input = Vec::with_capacity(13);
    crc_input.extend_from_slice(b"pHYs");
    crc_input.extend_from_slice(&chunk_data);

    let crc = crc32_png(&crc_input);

    let mut phys_chunk = Vec::with_capacity(21);
    phys_chunk.extend_from_slice(&9u32.to_be_bytes()); // length
    phys_chunk.extend_from_slice(b"pHYs");             // type
    phys_chunk.extend_from_slice(&chunk_data);          // data
    phys_chunk.extend_from_slice(&crc.to_be_bytes());   // CRC

    // Find the first IDAT chunk and insert pHYs before it.
    // PNG structure: 8-byte signature, then chunks (each: 4 len + 4 type + data + 4 crc)
    let mut result = Vec::with_capacity(png_data.len() + phys_chunk.len());
    let mut pos = 8; // skip PNG signature
    let mut phys_inserted = false;
    result.extend_from_slice(&png_data[..8]);

    while pos + 8 <= png_data.len() {
        let chunk_len = u32::from_be_bytes([
            png_data[pos], png_data[pos + 1], png_data[pos + 2], png_data[pos + 3],
        ]) as usize;
        let chunk_type = &png_data[pos + 4..pos + 8];
        // Use checked arithmetic to prevent overflow on malformed PNGs
        let Some(total_chunk_size) = chunk_len.checked_add(12) else {
            break; // chunk_len so large it overflows - malformed PNG
        };

        if pos.checked_add(total_chunk_size).map_or(true, |end| end > png_data.len()) {
            break;
        }

        // Skip any existing pHYs (we're replacing it)
        if chunk_type == b"pHYs" {
            pos += total_chunk_size;
            continue;
        }

        // Insert our pHYs before the first IDAT
        if chunk_type == b"IDAT" && !phys_inserted {
            result.extend_from_slice(&phys_chunk);
            phys_inserted = true;
        }

        result.extend_from_slice(&png_data[pos..pos + total_chunk_size]);
        pos += total_chunk_size;
    }

    // Copy any remaining bytes (e.g. trailing data after a truncated chunk)
    if pos < png_data.len() {
        result.extend_from_slice(&png_data[pos..]);
    }

    result
}

/// Simple CRC32 for PNG chunks (uses the standard PNG CRC polynomial)
fn crc32_png(data: &[u8]) -> u32 {
    let mut crc: u32 = 0xFFFF_FFFF;
    for &byte in data {
        crc ^= byte as u32;
        for _ in 0..8 {
            if crc & 1 != 0 {
                crc = (crc >> 1) ^ 0xEDB8_8320;
            } else {
                crc >>= 1;
            }
        }
    }
    crc ^ 0xFFFF_FFFF
}

impl Default for OutputEncoder {
    fn default() -> Self {
        Self::new()
    }
}