mozu-audio 0.2.0

//! Bounded microphone capture. Callback work is limited to conversion and lock-free enqueueing.
//! Resampling occurs after stop; cancelling drops the stream and its in-memory samples.
use crate::{Error, Result};
use cpal::{
    FromSample, Sample, SizedSample,
    traits::{DeviceTrait, HostTrait, StreamTrait},
};
use rtrb::{Consumer, Producer, RingBuffer};
use rubato::{FftFixedIn, Resampler};
use serde::Serialize;
use std::sync::{
    Arc,
    atomic::{AtomicU8, Ordering},
    mpsc::{self, Receiver, TryRecvError},
};

pub const SAMPLE_RATE: u32 = 16000;
pub const MIN_SAMPLES: usize = 10240;
pub const MAX_SAMPLES: usize = 491520;
const ACTIVE: u8 = 0;
const FULL: u8 = 1;
const INVALID_SAMPLE: u8 = 2;
const INCOMPLETE_FRAME: u8 = 3;

#[derive(Debug, Serialize)]
pub struct InputDevice {
    pub name: String,
    pub sample_rate: u32,
    pub channels: u16,
    pub sample_format: String,
}
fn device_description(
    device: &cpal::Device,
    config: &cpal::SupportedStreamConfig,
) -> Result<InputDevice> {
    Ok(InputDevice {
        name: device
            .name()
            .map_err(|e| Error::Invalid(format!("Read input device name: {e}")))?,
        sample_rate: config.sample_rate().0,
        channels: config.channels(),
        sample_format: format!("{:?}", config.sample_format()),
    })
}
pub fn default_input() -> Result<InputDevice> {
    let device = cpal::default_host()
        .default_input_device()
        .ok_or_else(|| Error::Invalid("No default microphone input device".into()))?;
    let config = device
        .default_input_config()
        .map_err(|e| Error::Invalid(format!("Read default input config: {e}")))?;
    device_description(&device, &config)
}

/// A single active take. Own this value to record; consume it to stop or cancel.
/// The queue is bounded by the selected capture limit, independently of UI polling.
pub struct Recording {
    stream: cpal::Stream,
    samples: Consumer<f32>,
    condition: Arc<AtomicU8>,
    errors: Receiver<cpal::StreamError>,
    device: InputDevice,
}
impl Recording {
    pub fn start() -> Result<Self> {
        Self::start_with_limit(MAX_SAMPLES)
    }

    /// Bound a single take by a mono 16 kHz sample count; live draining reuses the queue.
    pub fn start_with_limit(maximum_samples: usize) -> Result<Self> {
        if !(MIN_SAMPLES..=MAX_SAMPLES).contains(&maximum_samples) {
            return Err(Error::Invalid(format!(
                "Invalid capture limit: maximum_samples={maximum_samples}, supported={MIN_SAMPLES}..={MAX_SAMPLES}"
            )));
        }
        let device = cpal::default_host()
            .default_input_device()
            .ok_or_else(|| Error::Invalid("No default microphone input device".into()))?;
        let config = device
            .default_input_config()
            .map_err(|e| Error::Invalid(format!("Read default microphone config: {e}")))?;
        let description = device_description(&device, &config)?;
        if description.channels == 0
            || description.channels > 32
            || !(8000..=192000).contains(&description.sample_rate)
        {
            return Err(Error::Invalid(format!(
                "Unsupported microphone configuration: {description:?}"
            )));
        }
        let capacity = maximum_samples * description.sample_rate as usize / SAMPLE_RATE as usize;
        let (producer, samples) = RingBuffer::new(capacity);
        let condition = Arc::new(AtomicU8::new(ACTIVE));
        let (error_sender, errors) = mpsc::sync_channel(1);
        let channels = config.channels() as usize;
        let settings = config.config();
        let stream = match config.sample_format() {
            cpal::SampleFormat::F32 => build_stream::<f32>(
                &device,
                &settings,
                channels,
                producer,
                condition.clone(),
                error_sender,
            ),
            cpal::SampleFormat::I16 => build_stream::<i16>(
                &device,
                &settings,
                channels,
                producer,
                condition.clone(),
                error_sender,
            ),
            cpal::SampleFormat::I32 => build_stream::<i32>(
                &device,
                &settings,
                channels,
                producer,
                condition.clone(),
                error_sender,
            ),
            cpal::SampleFormat::U16 => build_stream::<u16>(
                &device,
                &settings,
                channels,
                producer,
                condition.clone(),
                error_sender,
            ),
            format => {
                return Err(Error::Invalid(format!(
                    "Unsupported microphone sample format: device={}, format={format:?}",
                    description.name
                )));
            }
        }
        .map_err(|e| {
            Error::Invalid(format!(
                "Build microphone stream: device={}, config={settings:?}, error={e}",
                description.name
            ))
        })?;
        stream.play().map_err(|e| {
            Error::Invalid(format!(
                "Start microphone stream: device={}, error={e}",
                description.name
            ))
        })?;
        Ok(Self {
            stream,
            samples,
            condition,
            errors,
            device: description,
        })
    }
    pub fn device(&self) -> &InputDevice {
        &self.device
    }
    pub fn limit_reached(&self) -> bool {
        self.condition.load(Ordering::Acquire) == FULL
    }
    pub fn check(&self) -> Result<()> {
        match self.errors.try_recv() {
            Ok(e) => {
                return Err(Error::Invalid(format!(
                    "Microphone stream failed: device={}, error={e}",
                    self.device.name
                )));
            }
            Err(TryRecvError::Empty | TryRecvError::Disconnected) => {}
        }
        check_condition(self.condition.load(Ordering::Acquire))
    }
    pub fn cancel(self) {
        drop(self);
    }
    /// Drain native-rate mono frames for live ASR. A full queue means input was
    /// lost and is an explicit error; callers must not pretend the session is complete.
    pub fn drain_input(&mut self, maximum: usize) -> Result<Vec<f32>> {
        self.check()?;
        if maximum == 0 || self.limit_reached() {
            return Err(Error::Invalid(
                "Live microphone overrun or zero drain size".into(),
            ));
        }
        let mut audio = Vec::with_capacity(maximum.min(self.samples.slots()));
        while audio.len() < maximum {
            match self.samples.pop() {
                Ok(value) => audio.push(value),
                Err(rtrb::PopError::Empty) => break,
            }
        }
        Ok(audio)
    }
    pub fn stop(self) -> Result<Vec<f32>> {
        let source_rate = self.device.sample_rate;
        to_model_rate(&self.stop_input()?, source_rate)
    }
    pub fn stop_input(self) -> Result<Vec<f32>> {
        self.check()?;
        let Self {
            stream,
            mut samples,
            condition,
            errors,
            device,
        } = self;
        drop(stream);
        if let Ok(e) = errors.try_recv() {
            return Err(Error::Invalid(format!(
                "Microphone stream failed at stop: device={}, error={e}",
                device.name
            )));
        }
        check_condition(condition.load(Ordering::Acquire))?;
        let mut mono = Vec::with_capacity(samples.slots());
        while let Ok(value) = samples.pop() {
            mono.push(value);
        }
        Ok(mono)
    }
}
fn check_condition(condition: u8) -> Result<()> {
    match condition {
        ACTIVE | FULL => Ok(()),
        INVALID_SAMPLE => Err(Error::Invalid(
            "Microphone supplied a non-finite sample".into(),
        )),
        INCOMPLETE_FRAME => Err(Error::Invalid(
            "Microphone callback supplied an incomplete channel frame".into(),
        )),
        value => Err(Error::Invalid(format!(
            "Invalid microphone condition: {value}"
        ))),
    }
}
fn build_stream<T>(
    device: &cpal::Device,
    config: &cpal::StreamConfig,
    channels: usize,
    mut producer: Producer<f32>,
    condition: Arc<AtomicU8>,
    error_sender: mpsc::SyncSender<cpal::StreamError>,
) -> std::result::Result<cpal::Stream, cpal::BuildStreamError>
where
    T: SizedSample,
    f32: FromSample<T>,
{
    device.build_input_stream(
        config,
        move |data: &[T], _| {
            if condition.load(Ordering::Acquire) != ACTIVE {
                return;
            }
            if !data.len().is_multiple_of(channels) {
                condition.store(INCOMPLETE_FRAME, Ordering::Release);
                return;
            }
            for frame in data.chunks_exact(channels) {
                let sample = crate::audio::mean_channels(
                    frame.iter().map(|&value| f32::from_sample(value)),
                    channels,
                );
                if !sample.is_finite() {
                    condition.store(INVALID_SAMPLE, Ordering::Release);
                    return;
                }
                if producer.push(sample).is_err() {
                    condition.store(FULL, Ordering::Release);
                    return;
                }
            }
        },
        move |error| {
            // Preserve the first failure; a full mailbox already holds the failure to report.
            match error_sender.try_send(error) {
                Ok(()) | Err(mpsc::TrySendError::Full(_)) => {}
                Err(mpsc::TrySendError::Disconnected(error)) => {
                    eprintln!("Microphone error after receiver closed: {error}")
                }
            }
        },
        None,
    )
}

/// Resample the complete mono take, including its tail, and remove FFT startup delay.
/// Rejects oversized input instead of silently truncating it.
pub fn to_model_rate(mono: &[f32], source_rate: u32) -> Result<Vec<f32>> {
    if !(8000..=192000).contains(&source_rate)
        || mono.is_empty()
        || mono.iter().any(|v| !v.is_finite())
    {
        return Err(Error::Invalid(format!(
            "Invalid recorded audio: source_rate={source_rate}, samples={}",
            mono.len()
        )));
    }
    let output_len = mono
        .len()
        .checked_mul(SAMPLE_RATE as usize)
        .ok_or_else(|| Error::Invalid("Audio sample count overflow".into()))?
        / source_rate as usize;
    if !(MIN_SAMPLES..=MAX_SAMPLES).contains(&output_len) {
        return Err(Error::Invalid(format!(
            "Record between 0.64 and 30.72 seconds: samples_16khz={output_len}, supported={MIN_SAMPLES}..={MAX_SAMPLES}"
        )));
    }
    if source_rate == SAMPLE_RATE {
        return Ok(mono.to_vec());
    }
    let mut resampler = make_resampler(source_rate)?;
    let delay = resampler.output_delay();
    let mut output = Vec::new();
    let mut consumed = 0;
    while output.len() < output_len + delay {
        let needed = resampler.input_frames_next();
        let end = (consumed + needed).min(mono.len());
        let channel = &mono[consumed..end];
        let channels = [channel];
        let input = if channel.is_empty() {
            None
        } else {
            Some(channels.as_slice())
        };
        let block = resampler.process_partial(input, None).map_err(|e| {
            Error::Invalid(format!(
                "Resample recording: source_rate={source_rate}, offset={consumed}, error={e}"
            ))
        })?;
        output.extend_from_slice(&block[0]);
        consumed = end;
    }
    Ok(output[delay..delay + output_len].to_vec())
}
/// Stateful native-rate to 16 kHz conversion. At most one second per push;
/// FFT history and fractional input are retained until finish flushes the tail.
pub struct RateConverter {
    source_rate: u32,
    resampler: Option<FftFixedIn<f32>>,
    pending: Vec<f32>,
    delay: usize,
    received: u64,
    emitted: u64,
    finished: bool,
}
fn make_resampler(source_rate: u32) -> Result<FftFixedIn<f32>> {
    let (mut gcd, mut remainder) = (source_rate as usize, SAMPLE_RATE as usize);
    while remainder != 0 {
        (gcd, remainder) = (remainder, gcd % remainder);
    }
    let unit = 2 * source_rate as usize / gcd;
    let chunk = (source_rate as usize / 50).div_ceil(unit) * unit;
    FftFixedIn::<f32>::new(source_rate as usize, SAMPLE_RATE as usize, chunk, 1, 1).map_err(|e| {
        Error::Invalid(format!(
            "Create resampler: source_rate={source_rate}, target_rate={SAMPLE_RATE}, error={e}"
        ))
    })
}
impl RateConverter {
    pub fn new(source_rate: u32) -> Result<Self> {
        if !(8000..=192000).contains(&source_rate) {
            return Err(Error::Invalid(format!(
                "Unsupported input rate: {source_rate}"
            )));
        }
        let resampler = if source_rate == SAMPLE_RATE {
            None
        } else {
            Some(make_resampler(source_rate)?)
        };
        let delay = resampler.as_ref().map_or(0, |r| r.output_delay());
        Ok(Self {
            source_rate,
            resampler,
            pending: Vec::new(),
            delay,
            received: 0,
            emitted: 0,
            finished: false,
        })
    }
    fn trim_delay(&mut self, block: &[f32]) -> Vec<f32> {
        let skip = self.delay.min(block.len());
        self.delay -= skip;
        let output = block[skip..].to_vec();
        self.emitted += output.len() as u64;
        output
    }
    pub fn push(&mut self, input: &[f32]) -> Result<Vec<f32>> {
        if self.finished {
            return Err(Error::Invalid("Resampler is finished".into()));
        }
        if input.len() > self.source_rate as usize || input.iter().any(|x| !x.is_finite()) {
            return Err(Error::Invalid(format!(
                "Resampler push must contain <=one second of finite samples: count={}, source_rate={}",
                input.len(),
                self.source_rate
            )));
        }
        self.received = self
            .received
            .checked_add(input.len() as u64)
            .ok_or_else(|| Error::Invalid("Resampler input clock overflow".into()))?;
        if self.resampler.is_none() {
            self.emitted += input.len() as u64;
            return Ok(input.to_vec());
        }
        self.pending.extend_from_slice(input);
        let mut output = Vec::new();
        loop {
            let resampler = self
                .resampler
                .as_mut()
                .ok_or_else(|| Error::Invalid("Missing rate converter".into()))?;
            let needed = resampler.input_frames_next();
            if self.pending.len() < needed {
                break;
            }
            let block = resampler
                .process(&[&self.pending[..needed]], None)
                .map_err(|e| {
                    Error::Invalid(format!(
                        "Streaming resampling failed: source_rate={}, error={e}",
                        self.source_rate
                    ))
                })?;
            self.pending.drain(..needed);
            output.extend(self.trim_delay(&block[0]));
        }
        Ok(output)
    }
    pub fn finish(&mut self) -> Result<Vec<f32>> {
        if self.finished {
            return Err(Error::Invalid("Resampler already finished".into()));
        }
        self.finished = true;
        let expected = self
            .received
            .checked_mul(SAMPLE_RATE as u64)
            .ok_or_else(|| Error::Invalid("Resampler output clock overflow".into()))?
            / self.source_rate as u64;
        if self.emitted > expected {
            return Err(Error::Invalid(
                "Resampler emitted beyond input duration".into(),
            ));
        }
        let remaining = (expected - self.emitted) as usize;
        let mut output = Vec::new();
        while output.len() < remaining {
            let channels = [self.pending.as_slice()];
            let input = if self.pending.is_empty() {
                None
            } else {
                Some(channels.as_slice())
            };
            let resampler = self.resampler.as_mut().ok_or_else(|| {
                Error::Invalid("Identity resampler has unexpected pending output".into())
            })?;
            let block = resampler.process_partial(input, None).map_err(|e| {
                Error::Invalid(format!(
                    "Flush resampler: source_rate={}, error={e}",
                    self.source_rate
                ))
            })?;
            self.pending.clear();
            output.extend(self.trim_delay(&block[0]));
        }
        output.truncate(remaining);
        self.emitted = expected;
        Ok(output)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn resampling_preserves_length_tone_and_last_syllable() {
        for rate in [8000, 22050, 44100, 48000, 96000] {
            let len = rate as usize + 173;
            let input: Vec<f32> = (0..len)
                .map(|i| (std::f32::consts::TAU * 1000.0 * i as f32 / rate as f32).sin() * 0.5)
                .collect();
            let output = to_model_rate(&input, rate).expect("valid take");
            assert_eq!(output.len(), len * 16000 / rate as usize);
            for range in [1000..2000, output.len() - 1000..output.len() - 100] {
                let error = range
                    .clone()
                    .map(|i| {
                        let expected =
                            (std::f32::consts::TAU * 1000.0 * i as f32 / 16000.0).sin() * 0.5;
                        (output[i] - expected).powi(2)
                    })
                    .sum::<f32>()
                    / range.len() as f32;
                assert!(error.sqrt() < 0.005, "rate={rate}, rms={}", error.sqrt());
            }
        }
    }
    #[test]
    fn resampling_flushes_tail_after_exact_input_blocks() {
        let input = vec![0.25; 48000];
        let output = to_model_rate(&input, 48000).expect("flush complete one-second recording");
        assert_eq!(output.len(), 16000);
        assert!((output[15900] - 0.25).abs() < 0.005);
    }
    #[test]
    fn streaming_resampler_matches_whole_clip_for_awkward_packets_and_tail() {
        for rate in [16000, 44100, 48000] {
            let audio: Vec<f32> = (0..rate as usize + 173)
                .map(|i| (i as f32 * 0.013).sin() * 0.2)
                .collect();
            let expected = to_model_rate(&audio, rate).unwrap();
            let mut converter = RateConverter::new(rate).unwrap();
            let mut actual = Vec::new();
            for packet in audio.chunks(137) {
                actual.extend(converter.push(packet).unwrap());
            }
            actual.extend(converter.finish().unwrap());
            assert_eq!(actual.len(), expected.len());
            assert!(
                actual
                    .iter()
                    .zip(expected)
                    .all(|(a, b)| (*a - b).abs() < 1e-6)
            );
            assert!(converter.finish().is_err());
            assert!(converter.push(&[0.0]).is_err());
        }
    }
    #[test]
    fn invalid_capture_limits_fail_before_opening_a_device() {
        for maximum in [0, MIN_SAMPLES - 1, MAX_SAMPLES + 1, usize::MAX] {
            assert!(matches!(
                Recording::start_with_limit(maximum),
                Err(Error::Invalid(_))
            ));
        }
    }

    #[test]
    fn resampling_rejects_invalid_or_oversized_takes() {
        assert!(to_model_rate(&[f32::NAN; 16000], 16000).is_err());
        assert!(to_model_rate(&[0.0; 100], 16000).is_err());
        assert!(to_model_rate(&vec![0.0; MAX_SAMPLES + 1], 16000).is_err());
    }
}