slimg-core 0.6.0

use std::ptr;

use libjxl_sys::*;

use crate::error::{Error, Result};

use super::types::EncodeConfig;

/// Safe wrapper around libjxl encoder.
pub(crate) struct Encoder {
    ptr: *mut JxlEncoder,
}

impl Encoder {
    /// Create a new JXL encoder instance.
    pub fn new() -> Result<Self> {
        let ptr = unsafe { JxlEncoderCreate(ptr::null()) };
        if ptr.is_null() {
            return Err(Error::Encode("failed to create JXL encoder".into()));
        }
        Ok(Self { ptr })
    }

    /// Encode interleaved 8-bit pixel data into JXL format.
    ///
    /// `num_channels` must be 3 (RGB) or 4 (RGBA); pass 3 for fully opaque
    /// images to avoid storing a redundant alpha channel.
    pub fn encode(
        &mut self,
        pixels: &[u8],
        width: u32,
        height: u32,
        num_channels: u32,
        config: &EncodeConfig,
    ) -> Result<Vec<u8>> {
        debug_assert!(num_channels == 3 || num_channels == 4);
        unsafe { JxlEncoderReset(self.ptr) };

        self.set_basic_info(width, height, num_channels, config)?;
        self.set_color_encoding()?;

        let frame_settings = unsafe { JxlEncoderFrameSettingsCreate(self.ptr, ptr::null()) };
        if frame_settings.is_null() {
            return Err(Error::Encode("failed to create frame settings".into()));
        }

        self.configure_frame(frame_settings, config)?;
        self.add_frame(frame_settings, pixels, width, height, num_channels)?;

        unsafe { JxlEncoderCloseInput(self.ptr) };

        self.process_output()
    }

    fn set_basic_info(
        &self,
        width: u32,
        height: u32,
        num_channels: u32,
        config: &EncodeConfig,
    ) -> Result<()> {
        let has_alpha = num_channels == 4;
        unsafe {
            // NOTE: libjxl recommends JxlEncoderInitBasicInfo; it is not in
            // the prebuilt bindings yet. Zeroed defaults are equivalent for
            // every field we do not set below (0 means "auto"/"none").
            let mut info: JxlBasicInfo = std::mem::zeroed();
            info.xsize = width;
            info.ysize = height;
            info.bits_per_sample = 8;
            info.exponent_bits_per_sample = 0;
            info.num_color_channels = 3;
            info.num_extra_channels = if has_alpha { 1 } else { 0 };
            info.alpha_bits = if has_alpha { 8 } else { 0 };
            info.alpha_exponent_bits = 0;
            info.orientation = JxlOrientation_JXL_ORIENT_IDENTITY;
            info.uses_original_profile = if config.lossless { 1 } else { 0 };

            check_status(JxlEncoderSetBasicInfo(self.ptr, &info), "set basic info")
        }
    }

    fn set_color_encoding(&self) -> Result<()> {
        unsafe {
            let mut color: JxlColorEncoding = std::mem::zeroed();
            JxlColorEncodingSetToSRGB(&mut color, 0); // is_gray = false
            check_status(
                JxlEncoderSetColorEncoding(self.ptr, &color),
                "set color encoding",
            )
        }
    }

    fn configure_frame(
        &self,
        settings: *mut JxlEncoderFrameSettings,
        config: &EncodeConfig,
    ) -> Result<()> {
        if config.lossless {
            unsafe {
                check_status(JxlEncoderSetFrameLossless(settings, 1), "set lossless")?;
            }
        }
        unsafe {
            check_status(
                JxlEncoderSetFrameDistance(settings, config.distance),
                "set distance",
            )
        }
    }

    fn add_frame(
        &self,
        settings: *mut JxlEncoderFrameSettings,
        pixels: &[u8],
        width: u32,
        height: u32,
        num_channels: u32,
    ) -> Result<()> {
        let format = JxlPixelFormat {
            num_channels,
            data_type: JxlDataType_JXL_TYPE_UINT8,
            endianness: JxlEndianness_JXL_NATIVE_ENDIAN,
            align: 0,
        };

        let expected = (width as usize) * (height as usize) * num_channels as usize;
        debug_assert_eq!(pixels.len(), expected);

        unsafe {
            check_status(
                JxlEncoderAddImageFrame(settings, &format, pixels.as_ptr().cast(), pixels.len()),
                "add image frame",
            )
        }
    }

    fn process_output(&self) -> Result<Vec<u8>> {
        let mut buffer = vec![0u8; 64 * 1024]; // 64 KB initial
        let mut all_output = Vec::new();

        loop {
            let mut next_out = buffer.as_mut_ptr();
            let mut avail_out = buffer.len();

            let status =
                unsafe { JxlEncoderProcessOutput(self.ptr, &mut next_out, &mut avail_out) };

            let written = buffer.len() - avail_out;
            all_output.extend_from_slice(&buffer[..written]);

            if status == JxlEncoderStatus_JXL_ENC_SUCCESS {
                return Ok(all_output);
            } else if status == JxlEncoderStatus_JXL_ENC_NEED_MORE_OUTPUT {
                continue;
            } else {
                return Err(Error::Encode("JXL encoding failed".into()));
            }
        }
    }
}

impl Drop for Encoder {
    fn drop(&mut self) {
        unsafe {
            JxlEncoderDestroy(self.ptr);
        }
    }
}

/// Check a `JxlEncoderStatus` and convert to `Result`.
///
/// # Safety
/// The caller must ensure `status` was returned by a valid libjxl encoder call.
unsafe fn check_status(status: JxlEncoderStatus, context: &str) -> Result<()> {
    if status == JxlEncoderStatus_JXL_ENC_SUCCESS {
        Ok(())
    } else {
        Err(Error::Encode(format!("jxl {context}: status {status}")))
    }
}