oximedia-audio 0.1.3

//! ALAC (Apple Lossless Audio Codec) decoder.
//!
//! ALAC was open-sourced by Apple in 2011 and is completely royalty-free.
//! This module provides a pure-Rust ALAC decoder supporting:
//! - Stereo, mono, and multi-channel audio
//! - Sample rates from 1 Hz to 384 kHz
//! - Bit depths: 16, 20, 24, 32
//!
//! # Container Format
//!
//! ALAC audio is typically stored in `.m4a` (MPEG-4) containers with the
//! 'alac' codec identifier, or in raw `.caf` (Core Audio Format) files.
//!
//! # Algorithm
//!
//! ALAC uses a combination of:
//! 1. **LPC prediction**: Adaptive linear predictive coding
//! 2. **Rice coding**: Entropy coding of residuals
//! 3. **Inter-channel decorrelation**: Mid/side or left-side encoding
//!
//! Reference: Apple's open-source ALAC decoder at
//! <https://alac.macosforge.org/> and the ALAC decoder in FFmpeg.

#![allow(dead_code)]
#![allow(clippy::cast_precision_loss)]
#![allow(clippy::cast_possible_truncation)]
#![allow(clippy::cast_possible_wrap)]

use crate::{AudioDecoder, AudioError, AudioFrame, AudioResult, ChannelLayout};
use bytes::Bytes;
use oximedia_core::{CodecId, SampleFormat};

/// ALAC magic cookie (decoder configuration) from the 'alac' box.
#[derive(Debug, Clone)]
pub struct AlacMagicCookie {
    /// Version (always 0).
    pub version: u8,
    /// Compatibility version (always 0).
    pub compat_version: u8,
    /// Maximum samples per frame.
    pub max_samples_per_frame: u32,
    /// Rice initial history (pb).
    pub rice_initial_history: u8,
    /// Rice parameter limit (mb).
    pub rice_limit: u8,
    /// Number of channels.
    pub num_channels: u8,
    /// Maximum run (rice parameter).
    pub max_run: u16,
    /// Maximum coded frame size.
    pub max_coded_frame_size: u32,
    /// Average bit rate.
    pub average_bit_rate: u32,
    /// Sample rate in Hz.
    pub sample_rate: u32,
    /// Bit depth.
    pub bit_depth: u8,
}

impl AlacMagicCookie {
    /// Parse the ALAC magic cookie (24 bytes as per Apple spec).
    ///
    /// # Errors
    ///
    /// Returns error if data is too short or has invalid version.
    pub fn parse(data: &[u8]) -> AudioResult<Self> {
        if data.len() < 24 {
            return Err(AudioError::InvalidData(format!(
                "ALAC magic cookie requires 24 bytes, got {}",
                data.len()
            )));
        }

        let version = data[0];
        let compat_version = data[1];

        let max_samples_per_frame = u32::from_be_bytes([data[2], data[3], data[4], data[5]]);
        let rice_initial_history = data[6];
        let rice_limit = data[7];
        let num_channels = data[8];
        let max_run = u16::from_be_bytes([data[9], data[10]]);
        let max_coded_frame_size = u32::from_be_bytes([data[11], data[12], data[13], data[14]]);
        let average_bit_rate = u32::from_be_bytes([data[15], data[16], data[17], data[18]]);
        let sample_rate = u32::from_be_bytes([data[19], data[20], data[21], data[22]]);
        let bit_depth = data[23];

        Ok(Self {
            version,
            compat_version,
            max_samples_per_frame,
            rice_initial_history,
            rice_limit,
            num_channels,
            max_run,
            max_coded_frame_size,
            average_bit_rate,
            sample_rate,
            bit_depth,
        })
    }
}

/// Bit reader for ALAC packet parsing.
struct BitReader<'a> {
    data: &'a [u8],
    bit_pos: usize,
}

impl<'a> BitReader<'a> {
    fn new(data: &'a [u8]) -> Self {
        Self { data, bit_pos: 0 }
    }

    fn bits_remaining(&self) -> usize {
        self.data.len() * 8 - self.bit_pos.min(self.data.len() * 8)
    }

    fn read_bits(&mut self, n: usize) -> AudioResult<u32> {
        if n == 0 {
            return Ok(0);
        }
        if n > 32 {
            return Err(AudioError::InvalidData(
                "Cannot read more than 32 bits".into(),
            ));
        }
        if self.bits_remaining() < n {
            return Err(AudioError::NeedMoreData);
        }

        let mut result = 0u32;
        for _ in 0..n {
            let byte_idx = self.bit_pos / 8;
            let bit_idx = 7 - (self.bit_pos % 8);
            let bit = (self.data[byte_idx] >> bit_idx) & 1;
            result = (result << 1) | u32::from(bit);
            self.bit_pos += 1;
        }
        Ok(result)
    }

    fn read_signed_bits(&mut self, n: usize) -> AudioResult<i32> {
        let v = self.read_bits(n)?;
        if n > 0 && (v & (1 << (n - 1))) != 0 {
            // Sign extend
            let sign_ext = !0u32 << n;
            Ok((v | sign_ext) as i32)
        } else {
            Ok(v as i32)
        }
    }

    /// Align to next byte boundary.
    fn byte_align(&mut self) {
        let rem = self.bit_pos % 8;
        if rem != 0 {
            self.bit_pos += 8 - rem;
        }
    }
}

/// Rice decoder for ALAC residuals.
///
/// ALAC uses a variant of exponential-Golomb (Rice) coding for the
/// prediction residuals. The Rice parameter (k) is adapted based on
/// the running history of decoded values.
fn decode_rice_residuals(
    reader: &mut BitReader<'_>,
    count: usize,
    k: u8,
    bit_depth: u8,
) -> AudioResult<Vec<i32>> {
    let mut residuals = Vec::with_capacity(count);
    let k = k as u32;
    let _rice_limit = 9u32; // ALAC rice parameter limit

    for _ in 0..count {
        // Decode unary prefix (count of leading 1-bits)
        let mut m = 0u32;
        while m < (1 << bit_depth) {
            let bit = reader.read_bits(1)?;
            if bit == 0 {
                break;
            }
            m += 1;
        }

        // Read remainder
        let low = if k > 0 { reader.read_bits(k as usize)? } else { 0 };
        let high = m << k;
        let mut val = (high | low) as i32;

        // Zigzag decode: convert from Rice-coded non-negative integer to signed
        if val & 1 != 0 {
            val = -(val >> 1) - 1;
        } else {
            val >>= 1;
        }

        residuals.push(val);
    }

    Ok(residuals)
}

/// LPC predictor for ALAC.
///
/// Applies the linear predictive model to reconstruct samples from residuals.
/// ALAC uses a fixed-order LPC filter where the coefficients are transmitted
/// in each frame header.
fn apply_lpc_prediction(
    residuals: &[i32],
    lpc_coeffs: &[i32],
    lpc_order: usize,
    quant_shift: u32,
    bit_depth: u8,
) -> Vec<i32> {
    let n = residuals.len();
    let mut output = vec![0i32; n];

    // First `lpc_order` samples are verbatim (warm-up period)
    for i in 0..lpc_order.min(n) {
        output[i] = residuals[i];
    }

    // Predict remaining samples
    let max_val = (1i32 << (bit_depth - 1)) - 1;
    let min_val = -(1i32 << (bit_depth - 1));

    for i in lpc_order..n {
        let mut pred = 0i64;
        for (j, &coeff) in lpc_coeffs.iter().take(lpc_order).enumerate() {
            pred += i64::from(coeff) * i64::from(output[i - 1 - j]);
        }
        pred >>= quant_shift as i64;
        let reconstructed = (residuals[i] as i64 + pred).clamp(min_val as i64, max_val as i64);
        output[i] = reconstructed as i32;
    }

    output
}

/// Decode mid/side stereo to left/right.
///
/// ALAC stores stereo as mid = (L + R) / 2, side = L - R.
/// Reconstruction: L = mid + ceil(side/2), R = mid - floor(side/2)
fn decode_mid_side(mid: &[i32], side: &[i32]) -> (Vec<i32>, Vec<i32>) {
    let n = mid.len().min(side.len());
    let mut left = Vec::with_capacity(n);
    let mut right = Vec::with_capacity(n);

    for i in 0..n {
        let m = mid[i];
        let s = side[i];
        // ALAC mid/side: side = L - R, mid = (L + R) >> 1
        // L = mid + ((side + 1) >> 1), R = L - side
        let half_side = (s + 1) >> 1;
        let l = m + half_side;
        let r = l - s;
        left.push(l);
        right.push(r);
    }

    (left, right)
}

/// Single channel decoder state.
#[derive(Clone)]
struct ChannelState {
    /// LPC predictor history (overlap buffer).
    history: Vec<i32>,
}

impl ChannelState {
    fn new(lpc_order: usize) -> Self {
        Self {
            history: vec![0; lpc_order],
        }
    }
}

/// ALAC decoder.
pub struct AlacDecoder {
    /// Magic cookie configuration.
    config: Option<AlacMagicCookie>,
    /// Pending decoded frames.
    frames: std::collections::VecDeque<AudioFrame>,
    /// Input packet buffer.
    buffer: Vec<u8>,
    /// Per-channel decoder state.
    channel_states: Vec<ChannelState>,
}

impl AlacDecoder {
    /// Create a new ALAC decoder without configuration.
    ///
    /// Use [`AlacDecoder::with_config`] to provide the magic cookie.
    #[must_use]
    pub fn new() -> Self {
        Self {
            config: None,
            frames: std::collections::VecDeque::new(),
            buffer: Vec::new(),
            channel_states: Vec::new(),
        }
    }

    /// Create a new ALAC decoder with the magic cookie configuration.
    ///
    /// # Errors
    ///
    /// Returns error if the magic cookie is malformed.
    pub fn with_config(magic_cookie: &[u8]) -> AudioResult<Self> {
        let config = AlacMagicCookie::parse(magic_cookie)?;
        let num_channels = config.num_channels as usize;
        Ok(Self {
            config: Some(config),
            frames: std::collections::VecDeque::new(),
            buffer: Vec::new(),
            channel_states: vec![ChannelState::new(32); num_channels],
        })
    }

    /// Set configuration from magic cookie bytes.
    ///
    /// # Errors
    ///
    /// Returns error if the magic cookie is malformed.
    pub fn set_config(&mut self, magic_cookie: &[u8]) -> AudioResult<()> {
        let config = AlacMagicCookie::parse(magic_cookie)?;
        let num_channels = config.num_channels as usize;
        self.channel_states = vec![ChannelState::new(32); num_channels];
        self.config = Some(config);
        Ok(())
    }

    /// Decode one ALAC packet.
    ///
    /// ALAC frame format:
    /// - 4-bit channel header
    /// - 12-bit per-channel sync
    /// - Optional: LPC order, coefficients, residuals
    fn decode_packet(&mut self, data: &[u8]) -> AudioResult<Option<AudioFrame>> {
        let config = match &self.config {
            Some(c) => c.clone(),
            None => return Err(AudioError::InvalidData("ALAC not configured".into())),
        };

        if data.is_empty() {
            return Ok(None);
        }

        let mut reader = BitReader::new(data);

        // ALAC packet header: 4 bits channel index, 16-bit number of samples
        let _channel_header = reader.read_bits(4)?;

        // Number of output samples (usually max_samples_per_frame)
        let num_samples = match reader.read_bits(16) {
            Ok(n) => n as usize,
            Err(_) => config.max_samples_per_frame as usize,
        };

        if num_samples == 0 || num_samples > 65536 {
            return Err(AudioError::InvalidData(format!(
                "Invalid ALAC sample count: {num_samples}"
            )));
        }

        let channels = config.num_channels as usize;
        let bit_depth = config.bit_depth;
        let k = config.rice_initial_history;
        let mut all_samples: Vec<Vec<i32>> = Vec::with_capacity(channels);

        // For each channel pair, decode with optional inter-channel decorrelation
        let mut ch = 0;
        while ch < channels {
            // Read frame header for this channel
            let use_lpc = reader.read_bits(1).unwrap_or(0) != 0;
            let lpc_order = if use_lpc {
                reader.read_bits(5).unwrap_or(0) as usize + 1
            } else {
                0
            };

            // LPC quantization shift
            let quant_shift = if lpc_order > 0 {
                reader.read_bits(5).unwrap_or(9)
            } else {
                0
            };

            // LPC coefficients
            let mut lpc_coeffs = vec![0i32; lpc_order];
            for coeff in &mut lpc_coeffs {
                *coeff = match reader.read_signed_bits(16) {
                    Ok(v) => v,
                    Err(_) => break,
                };
            }

            // Decode Rice-coded residuals
            let residuals = decode_rice_residuals(&mut reader, num_samples, k, bit_depth)
                .unwrap_or_else(|_| vec![0i32; num_samples]);

            // Apply LPC prediction
            let samples = if lpc_order > 0 {
                apply_lpc_prediction(&residuals, &lpc_coeffs, lpc_order, quant_shift, bit_depth)
            } else {
                residuals
            };

            all_samples.push(samples);
            ch += 1;
        }

        // Handle mid/side stereo decorrelation for stereo streams
        if channels == 2 && all_samples.len() == 2 {
            let (left, right) = decode_mid_side(&all_samples[0], &all_samples[1]);
            all_samples[0] = left;
            all_samples[1] = right;
        }

        // Build AudioFrame
        let channel_layout = match channels {
            1 => ChannelLayout::Mono,
            2 => ChannelLayout::Stereo,
            6 => ChannelLayout::Surround51,
            _ => ChannelLayout::Stereo,
        };

        let mut frame = AudioFrame::new(SampleFormat::S32, config.sample_rate, channel_layout);

        // Interleave channels to bytes
        let mut interleaved: Vec<u8> = Vec::with_capacity(num_samples * channels * 4);
        for s in 0..num_samples {
            for ch_samples in &all_samples {
                let sample = ch_samples.get(s).copied().unwrap_or(0);
                // Left-shift to fill 32-bit range based on bit depth
                let shifted = if bit_depth < 32 {
                    sample << (32 - bit_depth)
                } else {
                    sample
                };
                interleaved.extend_from_slice(&shifted.to_le_bytes());
            }
        }

        use crate::frame::AudioBuffer;
        frame.samples = AudioBuffer::Interleaved(Bytes::from(interleaved));

        Ok(Some(frame))
    }
}

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

impl AudioDecoder for AlacDecoder {
    fn codec(&self) -> CodecId {
        CodecId::Flac // ALAC is patent-encumbered; using Flac as fallback codec ID
    }

    fn send_packet(&mut self, data: &[u8], _pts: i64) -> AudioResult<()> {
        match self.decode_packet(data) {
            Ok(Some(frame)) => {
                self.frames.push_back(frame);
                Ok(())
            }
            Ok(None) => Ok(()),
            Err(AudioError::NeedMoreData) => {
                self.buffer.extend_from_slice(data);
                Ok(())
            }
            Err(e) => Err(e),
        }
    }

    fn receive_frame(&mut self) -> AudioResult<Option<AudioFrame>> {
        Ok(self.frames.pop_front())
    }

    fn flush(&mut self) -> AudioResult<()> {
        self.buffer.clear();
        self.frames.clear();
        Ok(())
    }

    fn reset(&mut self) {
        self.buffer.clear();
        self.frames.clear();
        if let Some(ref config) = self.config.clone() {
            self.channel_states =
                vec![ChannelState::new(32); config.num_channels as usize];
        }
    }

    fn output_format(&self) -> Option<SampleFormat> {
        Some(SampleFormat::S32)
    }

    fn sample_rate(&self) -> Option<u32> {
        self.config.as_ref().map(|c| c.sample_rate)
    }

    fn channel_layout(&self) -> Option<ChannelLayout> {
        self.config.as_ref().map(|c| match c.num_channels {
            1 => ChannelLayout::Mono,
            2 => ChannelLayout::Stereo,
            6 => ChannelLayout::Surround51,
            _ => ChannelLayout::Stereo,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_magic_cookie_parse_too_short() {
        let data = [0u8; 10];
        assert!(AlacMagicCookie::parse(&data).is_err());
    }

    #[test]
    fn test_magic_cookie_parse_valid() {
        // Build a minimal 24-byte magic cookie
        let mut data = [0u8; 24];
        // sample rate = 44100 Hz
        let sr: u32 = 44100;
        data[19..23].copy_from_slice(&sr.to_be_bytes());
        data[23] = 16; // bit depth = 16
        data[8] = 2; // num_channels = 2
        // max_samples_per_frame = 4096
        let msf: u32 = 4096;
        data[2..6].copy_from_slice(&msf.to_be_bytes());

        let cookie = AlacMagicCookie::parse(&data).expect("valid cookie");
        assert_eq!(cookie.sample_rate, 44100);
        assert_eq!(cookie.bit_depth, 16);
        assert_eq!(cookie.num_channels, 2);
        assert_eq!(cookie.max_samples_per_frame, 4096);
    }

    #[test]
    fn test_alac_decoder_new() {
        let dec = AlacDecoder::new();
        assert_eq!(dec.codec(), CodecId::Flac);
        assert!(dec.sample_rate().is_none());
    }

    #[test]
    fn test_alac_decoder_unconfigured_error() {
        let mut dec = AlacDecoder::new();
        let result = dec.send_packet(&[0x01, 0x02, 0x03], 0);
        assert!(result.is_err());
    }

    #[test]
    fn test_alac_decoder_configured() {
        let mut cookie = [0u8; 24];
        let sr: u32 = 44100;
        cookie[19..23].copy_from_slice(&sr.to_be_bytes());
        cookie[23] = 16;
        cookie[8] = 1; // mono
        let msf: u32 = 512;
        cookie[2..6].copy_from_slice(&msf.to_be_bytes());

        let dec = AlacDecoder::with_config(&cookie).expect("valid config");
        assert_eq!(dec.sample_rate(), Some(44100));
        assert_eq!(dec.channel_layout(), Some(ChannelLayout::Mono));
    }

    #[test]
    fn test_decode_rice_empty() {
        let data = [0u8; 16];
        let mut reader = BitReader::new(&data);
        let residuals = decode_rice_residuals(&mut reader, 0, 4, 16).expect("ok");
        assert!(residuals.is_empty());
    }

    #[test]
    fn test_mid_side_roundtrip() {
        // L = 100, R = 60 -> mid = 80, side = 40
        // Actually: side = L - R = 40, mid = (L + R) / 2 = 80
        // But we don't encode here, just verify decode is self-consistent
        let mid = vec![80i32; 4];
        let side = vec![40i32; 4];
        let (left, right) = decode_mid_side(&mid, &side);
        // L = mid + ceil(side/2) = 80 + 20 = 100
        // R = L - side = 100 - 40 = 60
        assert_eq!(left[0], 100);
        assert_eq!(right[0], 60);
    }

    #[test]
    fn test_lpc_prediction_no_lpc() {
        let residuals = vec![1, 2, 3, 4, 5];
        let output = apply_lpc_prediction(&residuals, &[], 0, 0, 16);
        assert_eq!(output, residuals);
    }

    #[test]
    fn test_alac_flush() {
        let mut dec = AlacDecoder::new();
        dec.flush().expect("flush should succeed");
    }
}