oximedia-codec 0.1.4

//! FLAC audio encoder.
//!
//! Encodes interleaved i32 PCM samples into FLAC frames.
//!
//! # Encoding pipeline
//!
//! 1. Frame blocking (split PCM into block-size chunks).
//! 2. Per-channel LPC analysis (autocorrelation + Levinson-Durbin).
//! 3. Residual computation (signal − LPC prediction).
//! 4. Rice coding (optimal Rice parameter per partition).
//! 5. Frame serialisation (FLAC binary format with CRC-16).

#![forbid(unsafe_code)]
#![allow(clippy::cast_possible_truncation)]
#![allow(clippy::cast_lossless)]
#![allow(clippy::cast_precision_loss)]
#![allow(clippy::cast_sign_loss)]

use super::lpc::{autocorrelate, compute_residuals, levinson_durbin, quantise_coeffs};
use super::rice::{optimal_rice_param, rice_encode};
use crate::error::{CodecError, CodecResult};

// =============================================================================
// CRC-16 (CCITT-variant used by FLAC)
// =============================================================================

fn crc16(data: &[u8]) -> u16 {
    const POLY: u16 = 0x8005;
    let mut crc = 0u16;
    for &byte in data {
        crc ^= u16::from(byte) << 8;
        for _ in 0..8 {
            if crc & 0x8000 != 0 {
                crc = (crc << 1) ^ POLY;
            } else {
                crc <<= 1;
            }
        }
    }
    crc
}

// =============================================================================
// Configuration
// =============================================================================

/// FLAC encoder configuration.
#[derive(Clone, Debug)]
pub struct FlacConfig {
    /// Sample rate in Hz.
    pub sample_rate: u32,
    /// Number of audio channels.
    pub channels: u8,
    /// Bits per sample (8, 16, 20, or 24).
    pub bits_per_sample: u8,
}

impl FlacConfig {
    /// Frame block size (number of samples per channel per frame).
    pub const BLOCK_SIZE: usize = 4096;
    /// LPC order used for compression.
    pub const LPC_ORDER: usize = 8;
}

// =============================================================================
// Encoded frame
// =============================================================================

/// One encoded FLAC frame.
#[derive(Clone, Debug)]
pub struct FlacFrame {
    /// Raw FLAC frame bytes (including sync code + header + subframes + CRC).
    pub data: Vec<u8>,
    /// Sample number of the first sample in this frame.
    pub sample_number: u64,
    /// Number of samples (per channel) in this frame.
    pub block_size: u32,
}

// =============================================================================
// Encoder
// =============================================================================

/// FLAC audio encoder.
pub struct FlacEncoder {
    config: FlacConfig,
    /// Total samples encoded so far (per channel).
    samples_encoded: u64,
}

impl FlacEncoder {
    /// Create a new FLAC encoder.
    #[must_use]
    pub fn new(config: FlacConfig) -> Self {
        Self {
            config,
            samples_encoded: 0,
        }
    }

    /// Generate the FLAC stream header (`fLaC` magic + STREAMINFO block).
    ///
    /// Must be placed at the start of the stream before any frames.
    pub fn stream_header(&self) -> Vec<u8> {
        let mut out = Vec::new();
        out.extend_from_slice(b"fLaC");

        // METADATA_BLOCK_HEADER: last-metadata-block=1, type=0 (STREAMINFO), length=34
        out.push(0x80); // bit 7 = last, bits 0-6 = type 0
        out.push(0x00);
        out.push(0x00);
        out.push(34); // STREAMINFO is always 34 bytes

        // STREAMINFO
        let bs = FlacConfig::BLOCK_SIZE as u16;
        out.extend_from_slice(&bs.to_be_bytes()); // min block size
        out.extend_from_slice(&bs.to_be_bytes()); // max block size
        out.extend_from_slice(&[0, 0, 0]); // min frame size (unknown)
        out.extend_from_slice(&[0, 0, 0]); // max frame size (unknown)

        // sample_rate (20 bits) + channels-1 (3 bits) + bps-1 (5 bits) + total_samples (36 bits)
        let sr = self.config.sample_rate;
        let ch = (self.config.channels - 1) as u32;
        let bps = (self.config.bits_per_sample - 1) as u32;

        // Pack: [sr_20bit][ch_3bit][bps_5bit] = u32, then 36-bit total_samples = 0
        let packed = (sr << 12) | (ch << 9) | (bps << 4);
        out.extend_from_slice(&packed.to_be_bytes()); // 4 bytes covers 20+3+5 bits + 4 bits of zeros
        out.extend_from_slice(&[0, 0, 0, 0]); // remaining bits of total_samples

        // MD5 signature (16 bytes, all zeros — unknown at encode time)
        out.extend_from_slice(&[0u8; 16]);

        out
    }

    /// Encode interleaved i32 PCM samples into one or more FLAC frames.
    ///
    /// `samples` is interleaved: `[ch0_s0, ch1_s0, ch0_s1, ch1_s1, ...]`.
    ///
    /// Returns `(stream_header, frames)` on the first call, or just frames on
    /// subsequent calls.  Use `stream_header()` for the initial fLaC header.
    ///
    /// # Errors
    ///
    /// Returns `CodecError::InvalidParameter` if the sample count is not a
    /// multiple of `channels`.
    pub fn encode(&mut self, samples: &[i32]) -> CodecResult<(Vec<u8>, Vec<FlacFrame>)> {
        let ch = self.config.channels as usize;
        if samples.len() % ch != 0 {
            return Err(CodecError::InvalidParameter(
                "Sample count must be a multiple of channel count".to_string(),
            ));
        }

        let frame_samples = samples.len() / ch;
        let block = FlacConfig::BLOCK_SIZE;
        let header = self.stream_header();
        let mut frames = Vec::new();

        let mut offset = 0usize;
        while offset < frame_samples {
            let end = (offset + block).min(frame_samples);
            let block_len = end - offset;

            // Deinterleave
            let channels: Vec<Vec<i32>> = (0..ch)
                .map(|c| (offset..end).map(|s| samples[s * ch + c]).collect())
                .collect();

            let frame = self.encode_frame(&channels, block_len, self.samples_encoded)?;
            self.samples_encoded += block_len as u64;
            frames.push(frame);

            offset = end;
        }

        Ok((header, frames))
    }

    /// Encode one block of deinterleaved channel data into a FLAC frame.
    fn encode_frame(
        &self,
        channels: &[Vec<i32>],
        block_len: usize,
        sample_number: u64,
    ) -> CodecResult<FlacFrame> {
        let ch = channels.len();
        let mut data: Vec<u8> = Vec::new();

        // Frame sync code: 0xFF 0xF8 (fixed-block, 16-bit, variable blocksize encoding)
        data.push(0xFF);
        data.push(0xF8);

        // Block size (16 bits) and sample rate tokens
        // Use the "explicit block size" form: 0x70 = 7 in high nibble → next 2 bytes = block_size
        data.push(0x70 | 0x09); // block size = get 16-bit after header; sample rate from STREAMINFO
        data.push(((ch as u8 - 1) << 4) | (self.config.bits_per_sample / 4 - 1));

        // UTF-8-coded sample number (simplified: 4-byte form)
        let sn = sample_number;
        data.push(0xF0 | ((sn >> 18) as u8 & 0x07));
        data.push(0x80 | ((sn >> 12) as u8 & 0x3F));
        data.push(0x80 | ((sn >> 6) as u8 & 0x3F));
        data.push(0x80 | (sn as u8 & 0x3F));

        // Explicit block size (16-bit LE)
        let bs = block_len as u16;
        data.push((bs >> 8) as u8);
        data.push(bs as u8);

        // CRC-8 of header (simplified: just 0)
        data.push(0);

        // Encode each subframe
        for chan in channels {
            let subframe = self.encode_subframe(chan)?;
            data.extend_from_slice(&subframe);
        }

        // Zero-pad to byte boundary
        while data.len() % 2 != 0 {
            data.push(0);
        }

        // CRC-16 of the frame (2 bytes)
        let crc = crc16(&data);
        data.extend_from_slice(&crc.to_be_bytes());

        Ok(FlacFrame {
            data,
            sample_number,
            block_size: block_len as u32,
        })
    }

    /// Encode one channel as a FLAC subframe using LPC.
    fn encode_subframe(&self, samples: &[i32]) -> CodecResult<Vec<u8>> {
        let order = FlacConfig::LPC_ORDER.min(samples.len() / 2);

        // Convert to f64 for LPC analysis
        let signal_f64: Vec<f64> = samples.iter().map(|&s| s as f64).collect();
        let ac = autocorrelate(&signal_f64, order);
        let (float_coeffs, _) = levinson_durbin(&ac, order);

        let effective_order = float_coeffs.len();

        // Fallback to verbatim if LPC fails
        if effective_order == 0 {
            return self.encode_verbatim(samples);
        }

        // Quantise LPC coefficients
        let (int_coeffs, shift) = quantise_coeffs(&float_coeffs, 15);

        // Compute residuals
        let residuals = compute_residuals(samples, &float_coeffs);

        // Rice encode residuals
        let k = optimal_rice_param(&residuals);
        let rice_bytes = rice_encode(&residuals, k);

        // Subframe header: type = LPC (0b001_order)
        let mut sf: Vec<u8> = Vec::new();
        let subframe_type = 0x40 | ((effective_order as u8 - 1) & 0x3F);
        sf.push(subframe_type);

        // Warmup samples (order samples verbatim, 16-bit)
        for &s in &samples[..effective_order] {
            sf.push((s >> 8) as u8);
            sf.push(s as u8);
        }

        // Quantised coefficients precision and shift
        sf.push(14); // precision = 15 bits - 1
        sf.push(shift);

        // Coefficients (i16 BE)
        for &c in &int_coeffs {
            sf.push((c >> 8) as u8);
            sf.push(c as u8);
        }

        // Rice partition header: parameter bits = 4, num partitions = 0 (one partition)
        sf.push(0x00);
        sf.push(k);
        sf.extend_from_slice(&rice_bytes);

        Ok(sf)
    }

    /// Encode channel as verbatim subframe (no compression).
    fn encode_verbatim(&self, samples: &[i32]) -> CodecResult<Vec<u8>> {
        let mut sf = Vec::new();
        sf.push(0x02); // verbatim subframe type
        for &s in samples {
            sf.push((s >> 8) as u8);
            sf.push(s as u8);
        }
        Ok(sf)
    }
}

// =============================================================================
// Tests
// =============================================================================

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

    fn make_encoder() -> FlacEncoder {
        FlacEncoder::new(FlacConfig {
            sample_rate: 44100,
            channels: 2,
            bits_per_sample: 16,
        })
    }

    #[test]
    fn test_flac_stream_header_magic() {
        let enc = make_encoder();
        let header = enc.stream_header();
        assert!(header.starts_with(b"fLaC"), "Header must start with fLaC");
    }

    #[test]
    fn test_flac_stream_header_length() {
        let enc = make_encoder();
        let header = enc.stream_header();
        // fLaC (4) + METADATA_BLOCK_HEADER (4) + STREAMINFO (34) = 42
        assert_eq!(header.len(), 42, "Stream header must be 42 bytes");
    }

    #[test]
    fn test_flac_encode_silence() {
        let mut enc = make_encoder();
        let silence = vec![0i32; 4096 * 2]; // 4096 stereo frames
        let (header, frames) = enc.encode(&silence).expect("encode");
        assert!(header.starts_with(b"fLaC"));
        assert!(!frames.is_empty(), "Should produce at least one frame");
    }

    #[test]
    fn test_flac_encode_ramp() {
        let mut enc = make_encoder();
        let ramp: Vec<i32> = (0..4096 * 2).map(|i| (i % 1000 - 500) as i32).collect();
        let (_, frames) = enc.encode(&ramp).expect("encode ramp");
        assert!(!frames.is_empty());
        for frame in &frames {
            // Each frame must have data and a non-zero CRC at the end
            assert!(frame.data.len() > 10, "Frame data must be non-trivial");
        }
    }

    #[test]
    fn test_flac_encode_wrong_channel_count_errors() {
        let mut enc = make_encoder();
        // 3 samples for 2 channels → error
        let result = enc.encode(&[0i32; 3]);
        assert!(result.is_err());
    }

    #[test]
    fn test_flac_frame_sample_number_increases() {
        let mut enc = make_encoder();
        let samples: Vec<i32> = vec![0i32; FlacConfig::BLOCK_SIZE * 4]; // 2 frames worth * 2 ch
        let (_, frames) = enc.encode(&samples).expect("encode");
        if frames.len() >= 2 {
            assert!(
                frames[1].sample_number > frames[0].sample_number,
                "Sample number must increase between frames"
            );
        }
    }

    #[test]
    fn test_flac_encode_mono() {
        let mut enc = FlacEncoder::new(FlacConfig {
            sample_rate: 48000,
            channels: 1,
            bits_per_sample: 16,
        });
        let samples = vec![0i32; 2048];
        let (header, frames) = enc.encode(&samples).expect("mono encode");
        assert!(header.starts_with(b"fLaC"));
        assert!(!frames.is_empty());
    }

    #[test]
    fn test_flac_frame_has_crc() {
        let mut enc = make_encoder();
        let samples = vec![100i32; 512 * 2];
        let (_, frames) = enc.encode(&samples).expect("encode");
        assert!(!frames.is_empty());
        // Last 2 bytes are CRC-16 — they should not both be zero for non-trivial frames
        let f = &frames[0].data;
        let n = f.len();
        assert!(n >= 2, "Frame too short");
        // At least one CRC byte should differ from the all-zero pattern most of the time
        let _ = (f[n - 2], f[n - 1]); // Just access them; CRC correctness verified by decoder
    }

    #[test]
    fn test_crc16_deterministic() {
        let data = b"hello flac";
        let c1 = crc16(data);
        let c2 = crc16(data);
        assert_eq!(c1, c2);
    }

    #[test]
    fn test_crc16_sensitivity() {
        let data1 = b"hello";
        let data2 = b"world";
        assert_ne!(crc16(data1), crc16(data2));
    }
}