pixstrip/pixstrip-core/src/encoder.rs

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

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

pub struct OutputEncoder;

impl OutputEncoder {
    pub fn new() -> Self {
        Self
    }

    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),
            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),
        }
    }

    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.webp_quality() as u8,
            ImageFormat::Png => preset.png_level(),
            _ => 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);

        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;
            let _ = started.write_scanlines(&pixels[start..end]);
        }

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

        Ok(output)
    }

    fn encode_png(&self, img: &image::DynamicImage) -> 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(),
        })?;

        let optimized = oxipng::optimize_from_memory(&buf, &oxipng::Options::default())
            .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 encoder = image::codecs::avif::AvifEncoder::new_with_speed_quality(
            cursor,
            6,
            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)
    }
}

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