avif-maker/src/encoder/avif.rs

use crate::app::EncodingSettings;
use crate::decoder::DecodedFrames;
use libavif::{AvifImage, Encoder, RgbPixels};
use std::fs::File;
use std::io::Write;
use std::path::Path;
use std::sync::Arc;
use thiserror::Error;

#[derive(Error, Debug)]
pub enum EncodeError {
    #[error("IO error: {0}")]
    Io(#[from] std::io::Error),
    #[error("AVIF encoding error: {0}")]
    Avif(String),
}

/// Progress callback type - receives (current, total)
pub type EncodeProgressCallback = Arc<dyn Fn(usize, usize) + Send + Sync>;

pub struct AvifEncoder {
    quality: u8,
    alpha_quality: u8,
    speed: u8,
    lossless: bool,
}

impl AvifEncoder {
    pub fn new(settings: &EncodingSettings) -> Self {
        Self {
            quality: if settings.lossless {
                100
            } else {
                settings.quality
            },
            alpha_quality: if settings.lossless {
                100
            } else {
                settings.alpha_quality
            },
            speed: settings.speed,
            lossless: settings.lossless,
        }
    }

    pub fn encode(&self, frames: &DecodedFrames, output_path: &Path) -> Result<(), EncodeError> {
        self.encode_with_progress(frames, output_path, None)
    }

    pub fn encode_with_progress(
        &self,
        frames: &DecodedFrames,
        output_path: &Path,
        progress: Option<EncodeProgressCallback>,
    ) -> Result<(), EncodeError> {
        if frames.frames.is_empty() {
            return Err(EncodeError::Avif("No frames to encode".to_string()));
        }

        let data = if frames.is_animated() {
            self.encode_animated_with_progress(frames, progress)?
        } else {
            // Report progress for single frame
            if let Some(ref cb) = progress {
                cb(0, 1);
            }
            let result = self.encode_single(frames)?;
            if let Some(ref cb) = progress {
                cb(1, 1);
            }
            result
        };

        let mut file = File::create(output_path)?;
        file.write_all(&data)?;

        Ok(())
    }

    fn encode_single(&self, frames: &DecodedFrames) -> Result<Vec<u8>, EncodeError> {
        let frame = &frames.frames[0];

        // Create AvifImage from RGBA data
        let image = self.create_avif_image(
            frame.width,
            frame.height,
            &frame.rgba,
            frames.has_alpha,
        )?;

        // Configure encoder
        let mut encoder = Encoder::new();
        self.configure_encoder(&mut encoder);

        // Encode single image
        let data = encoder
            .encode(&image)
            .map_err(|e| EncodeError::Avif(format!("Encoding failed: {:?}", e)))?;

        Ok(data.to_vec())
    }

    fn encode_animated(&self, frames: &DecodedFrames) -> Result<Vec<u8>, EncodeError> {
        self.encode_animated_with_progress(frames, None)
    }

    fn encode_animated_with_progress(
        &self,
        frames: &DecodedFrames,
        progress: Option<EncodeProgressCallback>,
    ) -> Result<Vec<u8>, EncodeError> {
        let mut encoder = Encoder::new();
        self.configure_encoder(&mut encoder);

        // Set timescale to milliseconds
        encoder.set_timescale(1000);

        let total_frames = frames.frames.len();

        for (i, frame) in frames.frames.iter().enumerate() {
            let image = self.create_avif_image(
                frame.width,
                frame.height,
                &frame.rgba,
                frames.has_alpha,
            )?;

            // Add frame with duration
            encoder
                .add_image(&image, frame.duration_ms, libavif::AddImageFlags::NONE)
                .map_err(|e| EncodeError::Avif(format!("Failed to add frame: {:?}", e)))?;

            // Report progress after each frame is added
            if let Some(ref cb) = progress {
                cb(i + 1, total_frames);
            }
        }

        let data = encoder
            .finish()
            .map_err(|e| EncodeError::Avif(format!("Failed to finish encoding: {:?}", e)))?;

        Ok(data.to_vec())
    }

    fn create_avif_image(
        &self,
        width: u32,
        height: u32,
        rgba: &[u8],
        has_alpha: bool,
    ) -> Result<AvifImage, EncodeError> {
        // Create RGB pixels from RGBA data
        let rgb = RgbPixels::new(width, height, rgba)
            .map_err(|e| EncodeError::Avif(format!("Failed to create RGB pixels: {:?}", e)))?;

        // Convert to AvifImage
        let image = rgb.to_image(if has_alpha {
            libavif::YuvFormat::Yuv444
        } else {
            libavif::YuvFormat::Yuv420
        });

        Ok(image)
    }

    fn configure_encoder(&self, encoder: &mut Encoder) {
        // Quality 0-100 maps directly to libavif quality
        encoder.set_quality(self.quality);
        encoder.set_alpha_quality(self.alpha_quality);
        encoder.set_speed(self.speed);

        if self.lossless {
            encoder.set_quality(100);
            encoder.set_alpha_quality(100);
        }
    }
}