chanora/poc/audio-capture-playback-spike/src/capture.rs

//! Audio capture: open the default input device, record N seconds of
//! 16-bit mono PCM, write to a WAV file.

use std::path::Path;
use std::sync::{Arc, Mutex};
use std::time::Duration;

use cpal::traits::{DeviceTrait, HostTrait, StreamTrait};
use cpal::{Sample, SampleFormat};
use hound::{SampleFormat as HoundSampleFormat, WavSpec, WavWriter};
use thiserror::Error;
use tracing::{info, warn};

#[derive(Debug, Error)]
pub enum AudioError {
    #[error("no input device available")]
    NoInputDevice,
    #[error("no output device available")]
    NoOutputDevice,
    #[error("device default config error: {0}")]
    DefaultConfig(#[from] cpal::DefaultStreamConfigError),
    #[error("device supported config error: {0}")]
    SupportedConfig(#[from] cpal::SupportedStreamConfigsError),
    #[error("stream build error: {0}")]
    BuildStream(#[from] cpal::BuildStreamError),
    #[error("stream play error: {0}")]
    PlayStream(#[from] cpal::PlayStreamError),
    #[error("wav io: {0}")]
    Wav(#[from] hound::Error),
    #[error("io: {0}")]
    Io(#[from] std::io::Error),
    #[error("unsupported sample format: {0:?}")]
    UnsupportedFormat(SampleFormat),
}

pub struct AudioCapture;

impl AudioCapture {
    /// Capture `duration` from the default input device into `path` as
    /// 16-bit mono PCM WAV. Returns the number of frames written.
    pub fn record_to_wav(path: &Path, duration: Duration) -> Result<u64, AudioError> {
        let host = cpal::default_host();
        let device = host.default_input_device().ok_or(AudioError::NoInputDevice)?;
        let device_name = device.name().unwrap_or_else(|_| "<unknown>".to_string());
        let config = device.default_input_config()?;
        let sample_rate = config.sample_rate().0;
        let channels = config.channels();
        let sample_format = config.sample_format();
        info!(
            target: "spike",
            device = %device_name,
            sample_rate, channels, ?sample_format,
            "opening capture stream"
        );

        let spec = WavSpec {
            channels: 1, // we down-mix to mono
            sample_rate,
            bits_per_sample: 16,
            sample_format: HoundSampleFormat::Int,
        };
        let writer = Arc::new(Mutex::new(Some(WavWriter::create(path, spec)?)));
        let frames = Arc::new(Mutex::new(0u64));

        let err_fn = |e| warn!(target: "spike", error = %e, "capture stream error");

        let stream = match sample_format {
            SampleFormat::F32 => build_input_stream::<f32>(
                &device,
                &config.into(),
                writer.clone(),
                frames.clone(),
                channels,
                err_fn,
            )?,
            SampleFormat::I16 => build_input_stream::<i16>(
                &device,
                &config.into(),
                writer.clone(),
                frames.clone(),
                channels,
                err_fn,
            )?,
            SampleFormat::U16 => build_input_stream::<u16>(
                &device,
                &config.into(),
                writer.clone(),
                frames.clone(),
                channels,
                err_fn,
            )?,
            other => return Err(AudioError::UnsupportedFormat(other)),
        };

        stream.play()?;
        std::thread::sleep(duration);
        drop(stream); // stops capture

        // Finalize the WAV.
        if let Some(w) = writer.lock().unwrap().take() {
            w.finalize()?;
        }
        let n = *frames.lock().unwrap();
        info!(target: "spike", frames = n, "capture finished");
        Ok(n)
    }
}

fn build_input_stream<T>(
    device: &cpal::Device,
    config: &cpal::StreamConfig,
    writer: Arc<Mutex<Option<WavWriter<std::io::BufWriter<std::fs::File>>>>>,
    frames: Arc<Mutex<u64>>,
    channels: u16,
    err_fn: fn(cpal::StreamError),
) -> Result<cpal::Stream, AudioError>
where
    T: SizedSample + ToI16,
{
    let stream = device.build_input_stream(
        config,
        move |data: &[T], _| {
            let mut w_guard = writer.lock().unwrap();
            if let Some(w) = w_guard.as_mut() {
                let mut n = 0u64;
                for frame in data.chunks(channels as usize) {
                    // Down-mix to mono by averaging channels.
                    let mut acc: i32 = 0;
                    for s in frame {
                        acc += s.to_i16() as i32;
                    }
                    let mono = (acc / frame.len() as i32) as i16;
                    let _ = w.write_sample(mono);
                    n += 1;
                }
                *frames.lock().unwrap() += n;
            }
        },
        err_fn,
        None,
    )?;
    Ok(stream)
}

/// `cpal::SizedSample` is the public marker. We expose it so the
/// generic stream builder can constrain `T` without leaking cpal types
/// out of this module.
pub trait SizedSample: cpal::SizedSample + Send + 'static {}
impl<T: cpal::SizedSample + Send + 'static> SizedSample for T {}

/// Adapter so we can convert any supported sample format down to i16
/// for WAV writing.
pub trait ToI16 {
    fn to_i16(&self) -> i16;
}

impl ToI16 for i16 {
    fn to_i16(&self) -> i16 {
        *self
    }
}

impl ToI16 for u16 {
    fn to_i16(&self) -> i16 {
        (i32::from(*self) - i32::from(i16::MAX) - 1) as i16
    }
}

impl ToI16 for f32 {
    fn to_i16(&self) -> i16 {
        let v = (*self * f32::from(i16::MAX)).clamp(f32::from(i16::MIN), f32::from(i16::MAX));
        v as i16
    }
}

// `Sample` trait import is required for cpal's older method shapes; we
// keep an explicit no-op use so this stays consistent across versions.
#[allow(dead_code)]
fn _assert_sample<T: Sample>() {}