midas_audio/codecs/

adpcm.rs

use crate::codecs::tables::{ADPCM_PRED_DIFF_TABLE, ADPCM_STEP_SIZE_ADJ, ADPCM_STEP_SIZES};
use std::fmt;

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum AdpcmError {
    InvalidMagic,
    BufferTooSmall { expected: u16, actual: u16 },
    ChecksumMismatch { expected: u8, actual: u8 },
    ChannelMismatch { expected: u16, actual: u16 },
    UnsupportedSampleType(u8),
}

impl fmt::Display for AdpcmError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            AdpcmError::InvalidMagic => write!(f, "Invalid ADP magic numbers"),
            AdpcmError::BufferTooSmall { expected, actual } => {
                write!(
                    f,
                    "Input buffer too small. Expected {expected}, got {actual}",
                )
            }
            AdpcmError::ChecksumMismatch { expected, actual } => {
                write!(
                    f,
                    "ADP header checksum mismatch. Expected {expected}, got {actual}",
                )
            }
            AdpcmError::ChannelMismatch { expected, actual } => {
                write!(
                    f,
                    "ADP channel mismatch. Header reports {actual}, state expects {expected}",
                )
            }
            AdpcmError::UnsupportedSampleType(t) => write!(f, "Unsupported ADP sample type: {t}"),
        }
    }
}

impl std::error::Error for AdpcmError {}

/// Defines the encoding format and channel layout for audio data.
#[repr(u8)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SampleType {
    /// No sample type specified or invalid type.
    None = 0,
    /// 8-bit Unsigned PCM (Mono).
    /// Samples are stored as single bytes with an offset of 128 (0 is -128, 128 is 0, 255 is 127).
    PcmMono8 = 1,
    /// 16-bit Signed PCM (Mono).
    /// Samples are stored as two-byte little-endian integers.
    PcmMono16 = 2,
    /// 8-bit Unsigned PCM (Stereo).
    /// Samples alternate between Left and Right channels, each as a single byte offset by 128.
    PcmStereo8 = 3,
    /// 16-bit Signed PCM (Stereo).
    /// Samples alternate between Left and Right channels, each as two-byte little-endian integers.
    PcmStereo16 = 4,
    /// 4-bit ADPCM (Mono).
    /// Adaptive Differential Pulse Code Modulation. Stores the difference between samples
    /// using 4 bits per sample, packed two samples per byte.
    AdpcmMono = 5,
    /// 4-bit ADPCM (Stereo).
    /// Adaptive Differential Pulse Code Modulation for two channels. Stores differences
    /// using 4 bits per sample (1 byte contains one L and one R sample nibble).
    AdpcmStereo = 6,
}

impl From<u8> for SampleType {
    fn from(v: u8) -> Self {
        match v {
            1 => SampleType::PcmMono8,
            2 => SampleType::PcmMono16,
            3 => SampleType::PcmStereo8,
            4 => SampleType::PcmStereo16,
            5 => SampleType::AdpcmMono,
            6 => SampleType::AdpcmStereo,
            _ => SampleType::None,
        }
    }
}

/// Holds the persistent state for an audio channel during ADPCM encoding/decoding.
#[derive(Debug, Clone, Copy, PartialEq, Default)]
pub struct AdpcmChannelState {
    /// The predicted 16-bit sample value.
    pub pred_value: i32,
    /// The index into the step size table (0-88).
    pub step_size_index: i32,
}

use crate::types::Channels;

/// Manages the state for ADPCM and PCM encoding/decoding within a stream.
///
/// This state handles the framing, headers, and adaptive predictor values for ADPCM.
/// Although it is named `AdpcmState`, it supports all formats defined in [`SampleType`],
/// including 8-bit and 16-bit PCM.
pub struct AdpcmState {
    pub channel_layout: Channels,
    pub frame_length: usize,
    pub channels: Vec<AdpcmChannelState>,
    pub frame_bytes: usize,
    pub header_bytes: usize,
    pub header_size: usize,
    pub encoding_type: SampleType,
    pub sample_rate: u32,
    // We handle header buffer logic slightly differently in Rust to avoid raw pointer casting,
    // but the concept remains: we need to enforce frame alignment.
    pub current_header: Option<Vec<u8>>,
    pub low_nibble: bool,
    pub pending_byte: u8,
    pub current_channel: usize,
}

impl AdpcmState {
    #[must_use]
    pub fn new(channel_layout: Channels, frame_length: usize) -> Self {
        let channels = vec![AdpcmChannelState::default(); channel_layout.count()];
        let encoding_type = match channel_layout {
            Channels::Stereo => SampleType::AdpcmStereo,
            Channels::Mono => SampleType::AdpcmMono,
        };
        Self {
            channel_layout,
            frame_length,
            channels,
            frame_bytes: 0,
            header_bytes: 0,
            header_size: 0,
            current_header: None,
            low_nibble: false,
            pending_byte: 0,
            current_channel: 0,
            encoding_type,
            sample_rate: 22050,
        }
    }

    #[must_use]
    pub fn with_type(
        channel_layout: Channels,
        frame_length: usize,
        encoding_type: SampleType,
        sample_rate: u32,
    ) -> Self {
        let channels: Vec<AdpcmChannelState> =
            vec![AdpcmChannelState::default(); channel_layout.count()];
        Self {
            channel_layout,
            frame_length,
            channels,
            frame_bytes: 0,
            header_bytes: 0,
            header_size: 0,
            current_header: None,
            low_nibble: false,
            pending_byte: 0,
            current_channel: 0,
            encoding_type,
            sample_rate,
        }
    }

    #[must_use]
    pub fn get_header_size(channel_layout: Channels) -> usize {
        6 + channel_layout.count() * 3
    }

    pub fn encode(&mut self, input: &[i16], output: &mut [u8]) -> (usize, usize) {
        let mut input_used = 0;
        let mut output_written = 0;
        let mut input_idx = 0;
        let mut output_idx = 0;

        while output_idx < output.len() && input_idx < input.len() {
            // Header writing logic
            if self.frame_bytes == 0 {
                if self.current_header.is_none() {
                    // Build generic header (6 bytes)
                    let mut header = Vec::new();
                    header.push(0x17);
                    header.push(0x42);
                    header.extend_from_slice(&(self.frame_length as u16).to_le_bytes());
                    header.push(0); // sum placeholder (idx 4)
                    header.push(self.encoding_type as u8); // type (idx 5)

                    // Channel headers (3 bytes each)
                    for i in 0..self.channel_layout.count() {
                        header
                            .extend_from_slice(&(self.channels[i].pred_value as i16).to_le_bytes());
                        header.push(self.channels[i].step_size_index as u8);
                    }

                    // Calculate checksum over entire header
                    let mut sum: u8 = 0;
                    for (i, b) in header.iter().enumerate() {
                        if i != 4 {
                            sum = sum.wrapping_add(*b);
                        }
                    }
                    header[4] = sum;

                    self.current_header = Some(header);
                    self.header_bytes = self.current_header.as_ref().unwrap().len();
                }

                // Write header to output
                if let Some(ref header) = self.current_header {
                    let offset = header.len() - self.header_bytes;
                    let available = output.len() - output_idx;
                    let to_write = std::cmp::min(self.header_bytes, available);

                    output[output_idx..output_idx + to_write]
                        .copy_from_slice(&header[offset..offset + to_write]);

                    output_idx += to_write;
                    output_written += to_write;
                    self.header_bytes -= to_write;

                    if self.header_bytes == 0 {
                        self.current_header = None;
                        self.frame_bytes =
                            self.frame_length - Self::get_header_size(self.channel_layout);
                        self.low_nibble = false;
                        self.pending_byte = 0;
                        self.current_channel = 0;
                    }
                }

                if output_idx >= output.len() {
                    break;
                }
            }

            // Encoding logic
            let ch = self.current_channel;
            let sample = input[input_idx];

            // Check whether the decoder can proceed with this sample
            let can_fit = match self.encoding_type {
                SampleType::AdpcmMono | SampleType::AdpcmStereo => {
                    self.low_nibble || output_idx < output.len()
                }
                SampleType::PcmMono16 | SampleType::PcmStereo16 => output_idx + 1 < output.len(),
                SampleType::PcmMono8 | SampleType::PcmStereo8 => output_idx < output.len(),
                _ => true,
            };

            if !can_fit {
                break;
            }

            // Consume sample
            input_idx += 1;
            input_used += 2;

            match self.encoding_type {
                SampleType::AdpcmMono | SampleType::AdpcmStereo => {
                    let chan = &mut self.channels[ch];

                    let mut diff = i32::from(sample) - chan.pred_value;
                    let mut qdiff = 0;

                    if diff < 0 {
                        qdiff = 8;
                        diff = -diff;
                    }

                    let mut step_size = ADPCM_STEP_SIZES[chan.step_size_index as usize];
                    let mut pred_diff = step_size / 8;

                    if diff >= step_size {
                        qdiff |= 4;
                        diff -= step_size;
                        pred_diff += step_size;
                    }
                    step_size /= 2;
                    if diff >= step_size {
                        qdiff |= 2;
                        diff -= step_size;
                        pred_diff += step_size;
                    }
                    step_size /= 2;
                    if diff >= step_size {
                        qdiff |= 1;
                        pred_diff += step_size;
                    }

                    if (qdiff & 8) != 0 {
                        chan.pred_value -= pred_diff;
                    } else {
                        chan.pred_value += pred_diff;
                    }

                    // Clip
                    chan.pred_value = chan.pred_value.clamp(-32768, 32767);

                    chan.step_size_index += ADPCM_STEP_SIZE_ADJ[(qdiff & 7) as usize];
                    chan.step_size_index = chan.step_size_index.clamp(0, 88);

                    if self.low_nibble {
                        output[output_idx] = self.pending_byte | (qdiff as u8);
                        output_idx += 1;
                        output_written += 1;
                        self.frame_bytes -= 1;
                        self.low_nibble = false;
                    } else {
                        self.pending_byte = (qdiff as u8) << 4;
                        self.low_nibble = true;
                    }
                }
                SampleType::PcmMono16 | SampleType::PcmStereo16 => {
                    let bytes = sample.to_le_bytes();
                    output[output_idx] = bytes[0];
                    output[output_idx + 1] = bytes[1];
                    output_idx += 2;
                    output_written += 2;
                    self.frame_bytes -= 2;
                }
                SampleType::PcmMono8 | SampleType::PcmStereo8 => {
                    // Convert i16 to u8 (offset 128)
                    let unsigned = ((i32::from(sample) + 32768) >> 8) as u8;
                    output[output_idx] = unsigned;
                    output_idx += 1;
                    output_written += 1;
                    self.frame_bytes -= 1;
                }
                _ => {}
            }

            // Next channel
            self.current_channel += 1;
            if self.current_channel >= self.channel_layout.count() {
                self.current_channel = 0;
            }
        }

        (input_used, output_written)
    }

    pub fn decode(
        &mut self,
        input: &[u8],
        output: &mut [i16],
    ) -> Result<(usize, usize), AdpcmError> {
        let mut input_used = 0;
        let mut output_written = 0;
        let mut input_idx = 0;
        let mut output_idx = 0;

        while input_idx < input.len() && output_idx < output.len() {
            if self.frame_bytes == 0 {
                let h_size = Self::get_header_size(self.channel_layout);
                if input.len() - input_idx < h_size {
                    break;
                }

                // Check magic: 0x17 0x42
                if input[input_idx] != 0x17 || input[input_idx + 1] != 0x42 {
                    if input_used > 0 {
                        // Attempt to find next magic if there's enough data
                        let mut found = false;
                        let mut search_idx = input_idx + 1;
                        while search_idx + 1 < input.len() {
                            if input[search_idx] == 0x17 && input[search_idx + 1] == 0x42 {
                                input_idx = search_idx;
                                found = true;
                                break;
                            }
                            search_idx += 1;
                        }
                        if !found {
                            break;
                        }
                    } else {
                        return Err(AdpcmError::InvalidMagic);
                    }
                }

                // Header info
                let frame_len =
                    u16::from_le_bytes([input[input_idx + 2], input[input_idx + 3]]) as usize;
                let read_sum = input[input_idx + 4];
                let header_type = input[input_idx + 5];

                // Validate checksum (sum mod 256 of entire header, except byte 4)
                let mut calc_sum: u8 = 0;
                for i in 0..h_size {
                    if i != 4 {
                        calc_sum = calc_sum.wrapping_add(input[input_idx + i]);
                    }
                }

                if calc_sum != read_sum {
                    return Err(AdpcmError::ChecksumMismatch {
                        expected: calc_sum,
                        actual: read_sum,
                    });
                }

                let encoding_type = SampleType::from(header_type);
                if encoding_type == SampleType::None {
                    return Err(AdpcmError::UnsupportedSampleType(header_type));
                }

                self.frame_length = frame_len;
                self.encoding_type = encoding_type;
                // self.sample_rate = 22050; // REMOVED: Do not overwrite sample rate from header

                let expected_channels = self.channel_layout.count();
                let actual_channels = match encoding_type {
                    SampleType::PcmMono8 | SampleType::PcmMono16 | SampleType::AdpcmMono => 1,
                    SampleType::PcmStereo8 | SampleType::PcmStereo16 | SampleType::AdpcmStereo => 2,
                    _ => 0,
                };

                if actual_channels != expected_channels {
                    return Err(AdpcmError::ChannelMismatch {
                        expected: expected_channels as u16,
                        actual: actual_channels as u16,
                    });
                }

                // Channel headers start at index 6
                let mut ptr = input_idx + 6;
                for i in 0..self.channel_layout.count() {
                    let pred = i32::from(i16::from_le_bytes([input[ptr], input[ptr + 1]]));
                    let step = i32::from(input[ptr + 2]);
                    self.channels[i].pred_value = pred;
                    self.channels[i].step_size_index = step.clamp(0, 88);
                    ptr += 3;
                }

                input_idx += h_size;
                input_used += h_size;
                self.frame_bytes = self.frame_length - h_size;
                self.low_nibble = false;
                self.current_channel = 0;
            }

            let ch = self.current_channel;
            if output_idx >= output.len() || self.frame_bytes == 0 {
                break;
            }

            match self.encoding_type {
                SampleType::AdpcmMono | SampleType::AdpcmStereo => {
                    let byte = input[input_idx];

                    let qdiff = if self.low_nibble {
                        let val = byte & 0x0F;
                        input_idx += 1;
                        input_used += 1;
                        self.frame_bytes -= 1;
                        val
                    } else {
                        (byte >> 4) & 0x0F
                    };
                    self.low_nibble = !self.low_nibble;

                    let chan = &mut self.channels[ch];
                    let pred_diff = ADPCM_PRED_DIFF_TABLE[(usize::try_from(chan.step_size_index)
                        .unwrap_or(0)
                        * 16)
                        + usize::from(qdiff)];

                    chan.pred_value += pred_diff;
                    chan.pred_value = chan.pred_value.clamp(-32768, 32767);

                    output[output_idx] = i16::try_from(chan.pred_value).unwrap_or(0);
                    output_idx += 1;
                    output_written += 2; // bytes written

                    chan.step_size_index += ADPCM_STEP_SIZE_ADJ[(qdiff & 7) as usize];
                    chan.step_size_index = chan.step_size_index.clamp(0, 88);
                }
                SampleType::PcmMono16 | SampleType::PcmStereo16 => {
                    if input_idx + 1 < input.len() {
                        let val = i16::from_le_bytes([input[input_idx], input[input_idx + 1]]);
                        output[output_idx] = val;
                        output_idx += 1;
                        output_written += 2;
                        input_idx += 2;
                        input_used += 2;
                        self.frame_bytes -= 2;
                    } else {
                        break;
                    }
                }
                SampleType::PcmMono8 | SampleType::PcmStereo8 => {
                    if input_idx < input.len() {
                        let b = input[input_idx];
                        // Convert u8 (offset 128) back to i16
                        let val = i16::try_from((i32::from(b) - 128) << 8).unwrap_or(0);
                        output[output_idx] = val;
                        output_idx += 1;
                        output_written += 2;
                        input_idx += 1;
                        input_used += 1;
                        self.frame_bytes -= 1;
                    } else {
                        break;
                    }
                }
                _ => {}
            }

            // Next channel
            self.current_channel += 1;
            if self.current_channel >= self.channel_layout.count() {
                self.current_channel = 0;
            }
        }

        Ok((input_used, output_written))
    }
}