crime 0.6.1

//! Concurrent real-time interface for multimedia engines.
//!
//! This crate provides abstractions and utilities for working with audio streams in Rust,
//! typically server-side applications that provides real-time audio streaming backed by
//! machine-learning models via gRPC, WebSocket, etc.
//!
//! This crate supports:
//!   - Audio inputs as `f32` samples
//!   - Resampling audio streams to different sample rates
//!   - Encoding audio streams to various formats (PCM, WAV, MP3)
//!
//! The main entry point is [`AudioStream`].

use async_stream::stream;
#[cfg(feature = "pcm")]
use audio_codec_algorithms::encode_ulaw;
use audioadapter_buffers::direct::SequentialSliceOfVecs;
use futures::StreamExt;
use futures::stream::Stream;
#[cfg(any(feature = "pcm", feature = "wav"))]
use half::f16;
use rubato::{Fft, FixedSync, Resampler};
use std::fmt::Debug;
use std::pin::Pin;
#[cfg(any(feature = "ogg", feature = "webm"))]
mod opus;
#[cfg(feature = "webm")]
mod webm;
mod wsola;
use wsola::time_scale;

#[cfg(feature = "mp3")]
pub type Mp3BitRate = mp3lame_encoder::Bitrate;
#[cfg(feature = "mp3")]
pub type Mp3Quality = mp3lame_encoder::Quality;
#[cfg(feature = "mp3")]
const MP3_FLUSH_MIN_BUFFER_SIZE: usize = 7_200;

#[cfg(any(feature = "ogg", feature = "webm"))]
pub type OpusApplication = ::opus::Application;
#[cfg(any(feature = "ogg", feature = "webm"))]
pub type OpusBitrate = ::opus::Bitrate;

/// Represents a PCM encoding.
/// A PCM encoding can be linear or companding (only G.711 μ-law is supported).
#[cfg(feature = "pcm")]
#[derive(Clone, Debug)]
pub enum PcmEncoding {
    /// Linear PCM.
    LinearPcm(LinearPcmEncoding),
    /// G.711 μ-law.
    G711MuLaw,
}

#[cfg(feature = "pcm")]
impl Default for PcmEncoding {
    fn default() -> Self {
        Self::LinearPcm(LinearPcmEncoding::default())
    }
}

/// Represents a linear PCM encoding.
///
/// All values are in little-endian format.
/// Float values are clamped between [-1.0, 1.0].
#[cfg(any(feature = "pcm", feature = "wav"))]
#[derive(Clone, Debug)]
pub enum LinearPcmEncoding {
    /// IEEE 754 half-precision floating point.
    Float16,
    /// IEEE 754 single-precision floating point.
    Float32,
    /// 16-bit signed integer.
    Int16,
}

#[cfg(any(feature = "pcm", feature = "wav"))]
impl Default for LinearPcmEncoding {
    fn default() -> Self {
        Self::Int16
    }
}

#[cfg(feature = "ogg")]
#[derive(Clone, Debug)]
pub enum OggContainer {
    Opus(OpusApplication, OpusBitrate),
}

#[cfg(feature = "webm")]
#[derive(Clone, Debug)]
pub enum WebmContainer {
    Opus(OpusApplication, OpusBitrate),
}

#[derive(Clone)]
pub enum AudioFormat {
    /// Raw stream of PCM samples.
    #[cfg(feature = "pcm")]
    Pcm(PcmEncoding),
    /// Stream of Linear PCM samples with WAV header.
    #[cfg(feature = "wav")]
    Wav(LinearPcmEncoding),
    /// Stream of MP3 samples.
    #[cfg(feature = "mp3")]
    Mp3(Mp3BitRate, Mp3Quality),
    /// Stream of Opus samples in an Ogg container.
    #[cfg(feature = "ogg")]
    Ogg(OggContainer),
    /// Stream of Opus samples in a WebM container.
    #[cfg(feature = "webm")]
    Webm(WebmContainer),
}

#[cfg(feature = "pcm")]
impl Default for AudioFormat {
    fn default() -> Self {
        Self::Pcm(PcmEncoding::default())
    }
}

impl Debug for AudioFormat {
    #[allow(unreachable_patterns, unused_variables)]
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            #[cfg(feature = "pcm")]
            AudioFormat::Pcm(pcm_encoding) => write!(f, "Pcm({:?})", pcm_encoding),
            #[cfg(feature = "wav")]
            AudioFormat::Wav(linear_pcm_encoding) => write!(f, "Wav({:?})", linear_pcm_encoding),
            #[cfg(feature = "mp3")]
            AudioFormat::Mp3(bit_rate, quality) => write!(
                f,
                "Mp3({}kbps, quality={})",
                *bit_rate as u16,
                mp3_quality_str(*quality)
            ),
            #[cfg(feature = "ogg")]
            AudioFormat::Ogg(codec) => write!(f, "Ogg({:?})", codec),
            #[cfg(feature = "webm")]
            AudioFormat::Webm(codec) => write!(f, "WebM({:?})", codec),
            _ => unreachable!(),
        }
    }
}

#[cfg(feature = "mp3")]
fn mp3_quality_str(quality: Mp3Quality) -> &'static str {
    match quality {
        Mp3Quality::Best => "Best",
        Mp3Quality::SecondBest => "SecondBest",
        Mp3Quality::NearBest => "NearBest",
        Mp3Quality::VeryNice => "VeryNice",
        Mp3Quality::Nice => "Nice",
        Mp3Quality::Good => "Good",
        Mp3Quality::Decent => "Decent",
        Mp3Quality::Ok => "Ok",
        Mp3Quality::SecondWorst => "SecondWorst",
        Mp3Quality::Worst => "Worst",
    }
}

/// Represents an audio stream with a specific sample rate.
///
/// This struct wraps a stream of audio samples (`f32`) and associates it with a sample rate.
/// It provides methods for audio processing, such as resampling and encoding.
///
/// # Fields
///
/// * `stream` - The input audio samples as a pinned, boxed stream of `f32` values.
/// * `sample_rate` - The sample rate of the audio stream, in Hz.
pub struct AudioStream<'a> {
    stream: Pin<Box<dyn Stream<Item = f32> + Send + 'a>>,
    sample_rate: u32,
}

impl<'a> AudioStream<'a> {
    /// Creates a new `AudioStream` from the given sample rate and input stream.
    ///
    /// # Arguments
    ///
    /// * `sample_rate` - The sample rate of the input audio stream, in Hz.
    /// * `stream` - The input audio samples as a stream of 32-bit floating point values (`f32`).
    ///
    /// # Returns
    ///
    /// A new `AudioStream` instance wrapping the provided stream and sample rate.
    pub fn new(sample_rate: u32, stream: impl Stream<Item = f32> + Send + 'a) -> AudioStream<'a> {
        AudioStream {
            sample_rate,
            stream: Box::pin(stream),
        }
    }

    /// Commits the audio manipulation on the original audio stream.
    ///
    /// This method applies resampling and encoding to the audio stream according to the specified
    /// output sample rate and encoding format. It consumes the `AudioStream` and returns a stream
    /// of encoded audio byte chunks.
    ///
    /// # Arguments
    ///
    /// * `sample_rate` - The desired output sample rate in Hz.
    /// * `time_scale_factor` - The time scale factor to apply to the audio stream. A value above
    ///   1.0 slows down the audio, a value below 1.0 speeds it up.
    /// * `format` - The desired output audio format.
    ///
    /// # Returns
    ///
    /// A pinned, boxed stream of encoded audio byte chunks (`Vec<u8>`).
    pub async fn commit(
        self,
        sample_rate: u32,
        time_scale_factor: f32,
        format: AudioFormat,
    ) -> Pin<Box<dyn Stream<Item = Vec<u8>> + Send + 'a>> {
        let time_scaled_stream = if (time_scale_factor - 1.0).abs() > f32::EPSILON {
            time_scale(self.stream, time_scale_factor, self.sample_rate).await
        } else {
            self.stream
        };

        let resampled_stream = if self.sample_rate != sample_rate {
            resample(time_scaled_stream, self.sample_rate, sample_rate).await
        } else {
            time_scaled_stream
        };

        encode(resampled_stream, sample_rate, format).await
    }
}

async fn resample<'a>(
    samples: impl Stream<Item = f32> + Send + 'a,
    input_rate: u32,
    output_rate: u32,
) -> Pin<Box<dyn Stream<Item = f32> + Send + 'a>> {
    if input_rate == output_rate {
        return Box::pin(samples);
    }
    let chunk_size = (20 * input_rate / 1000) as usize;
    let mut resampler = Fft::<f32>::new(
        input_rate as usize,
        output_rate as usize,
        chunk_size,
        1,
        1,
        FixedSync::Both,
    )
    .expect("creating resampler");
    let output_delay = resampler.output_delay();
    let padding_frames = resampler.input_frames_next();
    let mut input = Box::pin(samples.chain(futures::stream::repeat(0.0).take(padding_frames)));
    let out_frames_max = resampler.output_frames_max();
    Box::pin(
        stream! {
          let mut in_data = vec![vec![0.0f32; resampler.input_frames_max()]; 1];
          let mut out_data = vec![vec![0.0f32; out_frames_max]; 1];
          loop {
            let frames_needed = resampler.input_frames_next();

            in_data[0].clear();
            for _ in 0..frames_needed {
              match input.next().await {
                Some(sample) => in_data[0].push(sample),
                None => break,
              }
            }
            if in_data[0].len() < frames_needed {
              // No more pending samples. Not going to process the remainder.
              break;
            }

            let in_adapter = SequentialSliceOfVecs::new(&in_data, 1, frames_needed).unwrap();
            let mut out_adapter = SequentialSliceOfVecs::new_mut(&mut out_data, 1, out_frames_max).unwrap();
            let (_in_frames, out_frames) = resampler.process_into_buffer(
              &in_adapter, &mut out_adapter, None).unwrap();
            for sample in out_data[0][..out_frames].iter() {
              yield *sample;
            }
          }
        }
        .skip(output_delay),
    )
}

#[allow(unused_variables)]
async fn encode<'a>(
    samples: impl Stream<Item = f32> + Send + 'a,
    sample_rate: u32,
    format: AudioFormat,
) -> Pin<Box<dyn Stream<Item = Vec<u8>> + Send + 'a>> {
    match format {
        #[cfg(feature = "pcm")]
        AudioFormat::Pcm(PcmEncoding::LinearPcm(pcm_encoding)) => {
            encode_as_linear_pcm(samples, pcm_encoding).await
        }
        #[cfg(feature = "pcm")]
        AudioFormat::Pcm(PcmEncoding::G711MuLaw) => encode_as_g711_mu_law(samples).await,
        #[cfg(feature = "wav")]
        AudioFormat::Wav(linear_pcm_encoding) => {
            encode_as_wav(samples, sample_rate, linear_pcm_encoding).await
        }
        #[cfg(feature = "mp3")]
        AudioFormat::Mp3(bit_rate, quality) => {
            encode_as_mp3(samples, sample_rate, bit_rate, quality).await
        }
        #[cfg(feature = "ogg")]
        AudioFormat::Ogg(OggContainer::Opus(application, bitrate)) => {
            opus::encode_opus_as_ogg(samples, sample_rate, application, bitrate).await
        }
        #[cfg(feature = "webm")]
        AudioFormat::Webm(WebmContainer::Opus(application, bitrate)) => {
            opus::encode_opus_as_webm(samples, sample_rate, application, bitrate).await
        }
    }
}

/// Encode f32 samples in [-1.0, 1.0] to 8-bit G.711 μ-law bytes.
///
/// The implementation follows ITU-T G.711 with μ=255, mapping linear PCM to μ-law.
/// Input f32 is first clamped to [-1.0, 1.0], scaled to i16 range, then encoded.
#[cfg(feature = "pcm")]
async fn encode_as_g711_mu_law<'a>(
    samples: impl Stream<Item = f32> + Send + 'a,
) -> Pin<Box<dyn Stream<Item = Vec<u8>> + Send + 'a>> {
    let samples = Box::pin(samples);
    let mut sample_chunks = samples.ready_chunks(512);
    Box::pin(stream! {
      while let Some(chunk) = sample_chunks.next().await {
        let mut buf = Vec::with_capacity(chunk.len());
        for sample in chunk {
          let s = sample.clamp(-1.0, 1.0);
          // Scale to 16-bit linear PCM range
          let pcm16 = (s * 32767.0).round() as i16;
          buf.push(encode_ulaw(pcm16));
        }
        yield buf;
      }
    })
}

#[cfg(any(feature = "pcm", feature = "wav"))]
async fn encode_as_linear_pcm<'a>(
    samples: impl Stream<Item = f32> + Send + 'a,
    bit_depth: LinearPcmEncoding,
) -> Pin<Box<dyn Stream<Item = Vec<u8>> + Send + 'a>> {
    let samples = Box::pin(samples);
    let mut sample_chunks = samples.ready_chunks(512);
    let bytes_per_sample = match bit_depth {
        LinearPcmEncoding::Float16 | LinearPcmEncoding::Int16 => 2,
        LinearPcmEncoding::Float32 => 4,
    };
    Box::pin(stream! {
        while let Some(chunk) = sample_chunks.next().await {
            let mut buf = Vec::with_capacity(chunk.len() * bytes_per_sample);
            for sample in chunk {
                let sample = sample.clamp(-1.0, 1.0);
                match bit_depth {
                    LinearPcmEncoding::Float16 => {
                        buf.extend_from_slice(&f16::from_f32(sample).to_le_bytes());
                    }
                    LinearPcmEncoding::Float32 => {
                        buf.extend_from_slice(&sample.to_le_bytes());
                    }
                    LinearPcmEncoding::Int16 => {
                        let s = (sample * 32767.0).round() as i16;
                        buf.extend_from_slice(&s.to_le_bytes());
                    }
                };
            }
            yield buf;
        }
    })
}

#[cfg(feature = "wav")]
fn make_wav_header(sample_rate: u32, linear_pcm_encoding: &LinearPcmEncoding) -> [u8; 44] {
    let num_channels = 1u16;
    let bits_per_sample = match linear_pcm_encoding {
        LinearPcmEncoding::Float16 | LinearPcmEncoding::Int16 => 16u16,
        LinearPcmEncoding::Float32 => 32u16,
    };
    let audio_format = match linear_pcm_encoding {
        LinearPcmEncoding::Int16 => 1u16,
        LinearPcmEncoding::Float16 | LinearPcmEncoding::Float32 => 3u16,
    };
    let byte_rate = sample_rate * num_channels as u32 * (bits_per_sample as u32 / 8);
    let block_align = num_channels * (bits_per_sample / 8);
    let data_chunk_size = 0xFFFF_FFFFu32; // Unknown length for streaming
    let riff_chunk_size = 0xFFFF_FFFFu32; // Unknown length for streaming

    let mut header = [0u8; 44];
    header[0..4].copy_from_slice(b"RIFF");
    header[4..8].copy_from_slice(&(riff_chunk_size).to_le_bytes());
    header[8..12].copy_from_slice(b"WAVE");
    header[12..16].copy_from_slice(b"fmt ");
    header[16..20].copy_from_slice(&0x10u32.to_le_bytes()); // Subchunk1Size
    header[20..22].copy_from_slice(&audio_format.to_le_bytes());
    header[22..24].copy_from_slice(&num_channels.to_le_bytes());
    header[24..28].copy_from_slice(&sample_rate.to_le_bytes());
    header[28..32].copy_from_slice(&byte_rate.to_le_bytes());
    header[32..34].copy_from_slice(&block_align.to_le_bytes());
    header[34..36].copy_from_slice(&bits_per_sample.to_le_bytes());
    header[36..40].copy_from_slice(b"data");
    header[40..44].copy_from_slice(&data_chunk_size.to_le_bytes());
    header
}

#[cfg(feature = "wav")]
async fn encode_as_wav<'a>(
    samples: impl Stream<Item = f32> + Send + 'a,
    sample_rate: u32,
    linear_pcm_encoding: LinearPcmEncoding,
) -> Pin<Box<dyn Stream<Item = Vec<u8>> + Send + 'a>> {
    Box::pin(stream! {
        let header = make_wav_header(sample_rate, &linear_pcm_encoding);
        let mut pcm_stream = encode_as_linear_pcm(samples, linear_pcm_encoding).await;
        // Prepend header to the first PCM chunk so the first yielded Vec
        // always contains audio data (avoids a spurious early TTFB).
        if let Some(first_chunk) = pcm_stream.next().await {
            let mut buf = Vec::with_capacity(header.len() + first_chunk.len());
            buf.extend_from_slice(&header);
            buf.extend_from_slice(&first_chunk);
            yield buf;
        } else {
            // No audio — yield header only.
            yield header.to_vec();
            return;
        }
        while let Some(chunk) = pcm_stream.next().await {
            yield chunk;
        }
    })
}

#[cfg(feature = "mp3")]
async fn encode_as_mp3<'a>(
    samples: impl Stream<Item = f32> + Send + 'a,
    sample_rate: u32,
    bit_rate: Mp3BitRate,
    quality: Mp3Quality,
) -> Pin<Box<dyn Stream<Item = Vec<u8>> + Send + 'a>> {
    let mut mp3_encoder = mp3lame_encoder::Builder::new().expect("Create LAME encoder");
    mp3_encoder.set_num_channels(1).expect("set channels");
    mp3_encoder.set_brate(bit_rate).expect("set bit_rate");
    mp3_encoder
        .set_sample_rate(sample_rate)
        .expect("set sample rate");
    mp3_encoder.set_quality(quality).expect("set quality");
    let mut mp3_encoder = mp3_encoder.build().expect("To initialize LAME encoder");

    let samples = Box::pin(samples);
    let mut sample_chunks = samples.ready_chunks(128);
    Box::pin(stream! {
      let mut mp3_out_buffer = Vec::new();

      while let Some(chunk) = sample_chunks.next().await {
        let input = mp3lame_encoder::MonoPcm(&chunk);

        let encode_capacity =
            mp3lame_encoder::max_required_buffer_size(input.0.len());
        mp3_out_buffer.reserve(encode_capacity);
        mp3_encoder.encode_to_vec(input, &mut mp3_out_buffer).expect("To encode");

        if !mp3_out_buffer.is_empty() {
          yield take_output_chunk(
              &mut mp3_out_buffer,
              MP3_FLUSH_MIN_BUFFER_SIZE,
          );
        }
      }
      mp3_out_buffer.reserve(MP3_FLUSH_MIN_BUFFER_SIZE);
      mp3_encoder
          .flush_to_vec::<mp3lame_encoder::FlushNoGap>(&mut mp3_out_buffer)
          .expect("to flush");
      if !mp3_out_buffer.is_empty() {
        yield mp3_out_buffer;
      }
    })
}

#[cfg(feature = "mp3")]
fn take_output_chunk(buffer: &mut Vec<u8>, min_capacity: usize) -> Vec<u8> {
    let replacement_capacity = buffer.capacity().max(min_capacity);
    std::mem::replace(buffer, Vec::with_capacity(replacement_capacity))
}

#[cfg(all(test, feature = "mp3"))]
mod tests {
    use super::{
        AudioFormat, AudioStream, MP3_FLUSH_MIN_BUFFER_SIZE, Mp3BitRate, Mp3Quality,
        take_output_chunk,
    };
    use futures::{StreamExt, executor::block_on, stream};

    const INPUT_SAMPLE_RATE: u32 = 24_000;

    fn sample_data(sample_count: usize) -> Vec<f32> {
        (0..sample_count)
            .map(|i| ((i as f32 * 0.137).sin() * 0.8).clamp(-1.0, 1.0))
            .collect()
    }

    #[test]
    fn take_output_chunk_preserves_capacity_for_flush() {
        let mut buffer = Vec::with_capacity(MP3_FLUSH_MIN_BUFFER_SIZE + 123);
        buffer.extend_from_slice(b"frame");

        let expected_capacity = buffer.capacity();
        let yielded = take_output_chunk(&mut buffer, MP3_FLUSH_MIN_BUFFER_SIZE);

        assert_eq!(yielded, b"frame");
        assert!(buffer.is_empty());
        assert!(buffer.capacity() >= expected_capacity);
        assert!(buffer.capacity() >= MP3_FLUSH_MIN_BUFFER_SIZE);
    }

    #[test]
    fn mp3_streaming_completes_after_intermediate_yields() {
        block_on(async {
            let chunks = AudioStream::new(INPUT_SAMPLE_RATE, stream::iter(sample_data(96_000)))
                .commit(
                    INPUT_SAMPLE_RATE,
                    1.0,
                    AudioFormat::Mp3(Mp3BitRate::Kbps192, Mp3Quality::Best),
                )
                .await
                .collect::<Vec<_>>()
                .await;

            assert!(
                chunks.len() > 1,
                "expected at least one chunk before the final flush"
            );
            assert!(
                chunks.iter().all(|chunk| !chunk.is_empty()),
                "expected all emitted MP3 chunks to be non-empty"
            );
            assert!(
                chunks.iter().map(Vec::len).sum::<usize>() > 0,
                "expected MP3 stream to produce bytes"
            );
        });
    }
}