own_assist/own_mcp/src/audio/recording.rs

use cpal::traits::{DeviceTrait, StreamTrait};
use cpal::{Device, Error, SampleFormat, SizedSample, Stream, SupportedStreamConfig};
use hound::WavSpec;
use log::{debug, error, info, trace};
use std::fmt::Debug;
use std::sync::{Arc, Mutex};
use std::thread::JoinHandle;
use std::time::Duration;
use std::{path, thread};
use thiserror::Error;

fn generic_sample_to_i16(sample: impl SizedSample + Into<f64>, format: SampleFormat) -> i16 {
    match format {
        SampleFormat::I8 => (sample.into() * 2.0) as i16,
        SampleFormat::I16 => sample.into() as i16,
        SampleFormat::I32 => (sample.into() / 2.0) as i16,
        SampleFormat::F32 => (sample.into() * (i16::MAX as f64)) as i16,
        _ => panic!("Unsupported sample format {:?}", format),
    }
}

pub type SampleArc = Arc<Mutex<Vec<i16>>>;

fn stream_callback<T: SizedSample + Debug + Into<f64>>(
    input: &[T],
    samples: SampleArc,
    sample_format: SampleFormat,
) {
    if let Ok(mut guard) = samples.lock() {
        for &sample in input {
            guard.push(generic_sample_to_i16(sample, sample_format));
        }
    }
}

fn build_stream<T: SizedSample + Debug + hound::Sample + Send + Into<f64>>(
    mic: &Device,
    device_config: SupportedStreamConfig,
    samples: SampleArc,
    sample_format: SampleFormat,
) -> Result<Stream, Error> {
    mic.build_input_stream(
        device_config.into(),
        move |input: &[T], _info| stream_callback::<T>(input, samples.clone(), sample_format),
        |e| {
            error!("a stream error occurred while trying to record: {:?}", e);
        },
        None,
    )
}

#[derive(Debug)]
pub struct RecordingHandler {
    pub samples: SampleArc,
    pub spec: WavSpec,
    thread_handle: JoinHandle<Result<(), RecordingError>>,
    should_stop: Arc<Mutex<bool>>,
}

impl RecordingHandler {
    pub fn stop_recording(self) -> Result<(), RecordingError> {
        info!("Stopping recording");
        *self.should_stop.lock().unwrap() = true;
        self.thread_handle.join().unwrap()
    }
}

#[derive(Error, Debug)]
pub enum RecordingError {
    #[error("unsupported sample format")]
    UnsupportedSampleFormat(SampleFormat),
    #[error(transparent)]
    StreamError(#[from] cpal::Error),
    #[error("thread poisoned")]
    ThreadPoison,
    #[error("recording was already stopped")]
    RecordingAlreadyStopped,
}

fn start_recording_blocking_with_parameters(
    mic: &Device,
    device_config: SupportedStreamConfig,
    sample_format: SampleFormat,
    samples: SampleArc,
    should_stop: Arc<Mutex<bool>>,
) -> Result<(), RecordingError> {
    let stream = match sample_format {
        SampleFormat::I8 => {
            build_stream::<i8>(mic, device_config, samples.clone(), sample_format)?
        }
        SampleFormat::I16 => {
            build_stream::<i16>(mic, device_config, samples.clone(), sample_format)?
        }
        SampleFormat::I32 => {
            build_stream::<i32>(mic, device_config, samples.clone(), sample_format)?
        }
        SampleFormat::F32 => {
            build_stream::<f32>(mic, device_config, samples.clone(), sample_format)?
        }
        sample_format => {
            return Err(RecordingError::UnsupportedSampleFormat(sample_format));
        }
    };

    stream.play()?;

    loop {
        thread::sleep(Duration::from_millis(100));
        if let Ok(guard) = should_stop.lock()
            && *guard
        {
            debug!("Stopping recording gracefully");
            break;
        }
    }

    drop(stream);

    Ok(())
}

pub fn start(
    microphone: Device,
    device_config: SupportedStreamConfig,
) -> Arc<Mutex<Option<RecordingHandler>>> {
    let sample_format = device_config.sample_format();

    let samples: SampleArc = Arc::new(Mutex::new(Vec::new()));
    let recording_should_stop = Arc::new(Mutex::new(false));

    info!("started recording");
    debug!("sample format: {:?}", sample_format);

    let thread_samples = samples.clone();
    let thread_recording_stop = recording_should_stop.clone();
    let thread_device_config = device_config;
    let record_thread_handle = thread::spawn(move || {
        start_recording_blocking_with_parameters(
            &microphone,
            thread_device_config,
            sample_format,
            thread_samples,
            thread_recording_stop,
        )
    });

    let spec = WavSpec {
        channels: device_config.channels(),
        sample_rate: device_config.sample_rate(),
        bits_per_sample: 16,
        sample_format: hound::SampleFormat::Int,
    };

    debug!("device spec: {:?}", spec);

    Arc::new(Mutex::new(Some(RecordingHandler {
        spec,
        samples,
        thread_handle: record_thread_handle,
        should_stop: recording_should_stop,
    })))
}

pub fn stop_and_take_data(
    handler: Arc<Mutex<Option<RecordingHandler>>>,
) -> Result<(Vec<i16>, WavSpec), RecordingError> {
    let handler = handler
        .lock()
        .map_err(|_| RecordingError::ThreadPoison)?
        .take()
        .ok_or(RecordingError::RecordingAlreadyStopped)?;

    let samples = handler.samples.clone();
    let spec = handler.spec;

    handler.stop_recording()?;

    Ok((
        samples
            .lock()
            .map_err(|_| RecordingError::ThreadPoison)?
            .to_vec(),
        spec,
    ))
}

pub fn samples_to_wav(
    samples: impl IntoIterator<Item = i16>,
    spec: &WavSpec,
    filepath: impl AsRef<path::Path>,
) -> Result<Vec<u8>, hound::Error> {
    let wav_bytes = std::io::Cursor::new(Vec::<u8>::new());

    trace!(
        "creating a wav file writer to {}",
        filepath.as_ref().display()
    );
    let mut file_writer = hound::WavWriter::create(&filepath, *spec)?;

    trace!("writing samples...");
    for sample in samples {
        file_writer.write_sample(sample)?;
    }

    trace!("finalizing for file {}", filepath.as_ref().display());
    file_writer.finalize()?;

    Ok(wav_bytes.into_inner())
}