tenant.openagents/omega

1use std::{
2    env,
3    num::NonZero,
4    path::{Path, PathBuf},
5    sync::{Arc, Mutex},
6    time::Duration,
7};
8
9use anyhow::{Context, Result};
10use cpal::DeviceId;
11use cpal::traits::{DeviceTrait, StreamTrait};
12use rodio::{buffer::SamplesBuffer, conversions::SampleTypeConverter};
13use util::ResultExt;
14
15pub struct CaptureInput {
16    pub name: String,
17    pub input_device: Option<DeviceId>,
18    config: cpal::SupportedStreamConfig,
19    samples: Arc<Mutex<Vec<i16>>>,
20    _stream: cpal::Stream,
21}
22
23impl CaptureInput {
24    pub fn start(input_device: Option<DeviceId>) -> anyhow::Result<Self> {
25        let (device, config) = crate::default_device(true, input_device.as_ref())?;
26        let name = device
27            .description()
28            .map(|desc| desc.name().to_string())
29            .unwrap_or("<unknown>".to_string());
30        log::info!("Using microphone: {}", name);
31
32        let samples = Arc::new(Mutex::new(Vec::new()));
33        let stream = start_capture(device, config.clone(), samples.clone())?;
34
35        Ok(Self {
36            name,
37            input_device,
38            _stream: stream,
39            config,
40            samples,
41        })
42    }
43
44    pub fn finish(self) -> Result<PathBuf> {
45        let name = self.name;
46        let mut path = env::current_dir().context("Could not get current dir")?;
47        path.push(&format!("test_recording_{name}.wav"));
48        log::info!("Test recording written to: {}", path.display());
49        write_out(self.samples, self.config, &path)?;
50        Ok(path)
51    }
52}
53
54fn start_capture(
55    device: cpal::Device,
56    config: cpal::SupportedStreamConfig,
57    samples: Arc<Mutex<Vec<i16>>>,
58) -> Result<cpal::Stream> {
59    let stream = device
60        .build_input_stream_raw(
61            &config.config(),
62            config.sample_format(),
63            move |data, _: &_| {
64                let data = crate::get_sample_data(config.sample_format(), data).log_err();
65                let Some(data) = data else {
66                    return;
67                };
68                samples
69                    .try_lock()
70                    .expect("Only locked after stream ends")
71                    .extend_from_slice(&data);
72            },
73            |err| log::error!("error capturing audio track: {:?}", err),
74            Some(Duration::from_millis(100)),
75        )
76        .context("failed to build input stream")?;
77
78    stream.play()?;
79    Ok(stream)
80}
81
82fn write_out(
83    samples: Arc<Mutex<Vec<i16>>>,
84    config: cpal::SupportedStreamConfig,
85    path: &Path,
86) -> Result<()> {
87    let samples = std::mem::take(
88        &mut *samples
89            .try_lock()
90            .expect("Stream has ended, callback cant hold the lock"),
91    );
92    let samples: Vec<f32> = SampleTypeConverter::<_, f32>::new(samples.into_iter()).collect();
93    let mut samples = SamplesBuffer::new(
94        NonZero::new(config.channels()).expect("config channel is never zero"),
95        NonZero::new(config.sample_rate()).expect("config sample_rate is never zero"),
96        samples,
97    );
98    match rodio::wav_to_file(&mut samples, path) {
99        Ok(_) => Ok(()),
100        Err(e) => Err(anyhow::anyhow!("Failed to write wav file: {}", e)),
101    }
102}
103