Skip to main content

murmur_core/audio/
capture.rs

1use anyhow::{Context, Result};
2use cpal::traits::{DeviceTrait, HostTrait, StreamTrait};
3use hound::{SampleFormat, WavSpec, WavWriter};
4use nnnoiseless::DenoiseState;
5use std::collections::VecDeque;
6use std::fs::File;
7use std::io::BufWriter;
8use std::path::{Path, PathBuf};
9use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
10use std::sync::{Arc, Mutex};
11
12/// Duration of the pre-roll buffer in milliseconds.
13/// Captures audio *before* the user presses record so the first word isn't clipped.
14const PRE_ROLL_MS: u32 = 200;
15
16/// Number of 16 kHz samples in the pre-roll buffer.
17const PRE_ROLL_SAMPLES: usize = (TARGET_RATE * PRE_ROLL_MS / 1000) as usize;
18
19/// Expected recording duration for initial buffer capacity (seconds).
20/// Pre-allocating ~5 seconds of 16 kHz audio (80 000 samples ≈ 320 KB)
21/// avoids multiple early reallocations during typical dictation.
22const INITIAL_RECORDING_SECS: usize = 5;
23
24/// Shared state between the audio callback thread and the main thread.
25struct SharedCaptureState {
26    /// True while actively recording (as opposed to standby pre-roll capture).
27    recording: AtomicBool,
28    /// WAV file writer, set only during file-backed recordings.
29    writer: Mutex<Option<WavWriter<BufWriter<File>>>>,
30    /// In-memory buffer of 16 kHz mono f32 samples captured since `start()`.
31    samples: Arc<Mutex<Vec<f32>>>,
32    /// Ring buffer capturing recent audio while in standby mode.
33    /// When recording starts, these samples are drained into `samples`
34    /// so the beginning of speech is preserved.
35    pre_roll: Mutex<VecDeque<f32>>,
36    /// Count of samples dropped due to lock contention.
37    dropped_samples: AtomicU64,
38    /// Monotonically increasing count of audio callback invocations.
39    /// Used to detect dead streams (e.g. when a Bluetooth device disconnects).
40    callback_count: AtomicU64,
41}
42
43impl SharedCaptureState {
44    fn new() -> Self {
45        let initial_capacity = TARGET_RATE as usize * INITIAL_RECORDING_SECS;
46        Self {
47            recording: AtomicBool::new(false),
48            writer: Mutex::new(None),
49            samples: Arc::new(Mutex::new(Vec::with_capacity(initial_capacity))),
50            pre_roll: Mutex::new(VecDeque::with_capacity(PRE_ROLL_SAMPLES + 512)),
51            dropped_samples: AtomicU64::new(0),
52            callback_count: AtomicU64::new(0),
53        }
54    }
55
56    /// Dispatch processed audio samples to the appropriate buffer.
57    /// Called from the audio callback after mixing/resampling/denoising.
58    fn dispatch_samples(&self, samples: &[f32]) {
59        self.callback_count.fetch_add(1, Ordering::Relaxed);
60        if self.recording.load(Ordering::Acquire) {
61            if let Ok(mut buf) = self.samples.try_lock() {
62                buf.extend_from_slice(samples);
63            } else {
64                self.dropped_samples
65                    .fetch_add(samples.len() as u64, Ordering::Relaxed);
66            }
67            if let Ok(mut guard) = self.writer.try_lock() {
68                if let Some(ref mut w) = *guard {
69                    for &sample in samples {
70                        let _ = w.write_sample(f32_to_i16(sample));
71                    }
72                }
73            }
74        } else if let Ok(mut ring) = self.pre_roll.try_lock() {
75            for &s in samples {
76                if ring.len() >= PRE_ROLL_SAMPLES {
77                    ring.pop_front();
78                }
79                ring.push_back(s);
80            }
81        }
82    }
83}
84
85/// nnnoiseless operates on 480-sample frames at 48 kHz.
86const DENOISE_FRAME_SIZE: usize = DenoiseState::FRAME_SIZE;
87
88/// nnnoiseless native sample rate.
89const DENOISE_RATE: u32 = 48_000;
90
91/// Holds the denoiser state and an accumulation buffer for incomplete frames.
92/// Created once per recording session and shared with the audio callback via Arc<Mutex>.
93struct Denoiser {
94    state: Box<DenoiseState<'static>>,
95    /// Accumulates 16 kHz samples until we have enough to fill a 48 kHz frame.
96    pending_16k: Vec<f32>,
97    /// Collects denoised 16 kHz output samples ready for the consumer.
98    output_16k: Vec<f32>,
99    /// Whether to skip the very first output frame (fade-in artifact).
100    first_frame: bool,
101    // Pre-allocated scratch buffers to avoid per-callback heap allocations.
102    chunk_buf: Vec<f32>,
103    upsampled_buf: Vec<f32>,
104    denoised_48k_buf: Vec<f32>,
105    downsampled_buf: Vec<f32>,
106}
107
108impl Denoiser {
109    fn new() -> Self {
110        let frame_16k = DENOISE_FRAME_SIZE / 3; // 160
111        Self {
112            state: DenoiseState::new(),
113            pending_16k: Vec::with_capacity(frame_16k + 16),
114            output_16k: Vec::new(),
115            first_frame: true,
116            chunk_buf: Vec::with_capacity(frame_16k),
117            upsampled_buf: Vec::with_capacity(DENOISE_FRAME_SIZE),
118            denoised_48k_buf: Vec::with_capacity(DENOISE_FRAME_SIZE),
119            downsampled_buf: Vec::with_capacity(frame_16k),
120        }
121    }
122
123    /// Reset all state so the denoiser is clean for a new recording session.
124    fn reset(&mut self) {
125        self.pending_16k.clear();
126        self.output_16k.clear();
127        self.first_frame = true;
128        self.state = DenoiseState::new();
129        // Scratch buffers keep their allocations; just clear contents.
130        self.chunk_buf.clear();
131        self.upsampled_buf.clear();
132        self.denoised_48k_buf.clear();
133        self.downsampled_buf.clear();
134    }
135
136    /// Feed 16 kHz samples and return denoised 16 kHz samples.
137    ///
138    /// Internally upsamples to 48 kHz, runs nnnoiseless frame-by-frame,
139    /// then downsamples back to 16 kHz.
140    fn process(&mut self, samples_16k: &[f32]) -> &[f32] {
141        self.output_16k.clear();
142        self.pending_16k.extend_from_slice(samples_16k);
143
144        // Each 48 kHz frame of 480 samples corresponds to 160 samples at 16 kHz.
145        let frame_16k = DENOISE_FRAME_SIZE / 3; // 160
146
147        while self.pending_16k.len() >= frame_16k {
148            self.chunk_buf.clear();
149            self.chunk_buf.extend(self.pending_16k.drain(..frame_16k));
150
151            resample_into(
152                &self.chunk_buf,
153                TARGET_RATE,
154                DENOISE_RATE,
155                &mut self.upsampled_buf,
156            );
157
158            // nnnoiseless expects f32 in i16 range
159            let mut input_frame = [0.0f32; DENOISE_FRAME_SIZE];
160            for (i, &s) in self
161                .upsampled_buf
162                .iter()
163                .take(DENOISE_FRAME_SIZE)
164                .enumerate()
165            {
166                input_frame[i] = s * 32767.0;
167            }
168
169            let mut output_frame = [0.0f32; DENOISE_FRAME_SIZE];
170            self.state.process_frame(&mut output_frame, &input_frame);
171
172            if self.first_frame {
173                self.first_frame = false;
174                continue;
175            }
176
177            // Convert back from i16 range to [-1, 1]
178            self.denoised_48k_buf.clear();
179            self.denoised_48k_buf.extend(
180                output_frame
181                    .iter()
182                    .map(|&s| (s / 32767.0_f32).clamp(-1.0, 1.0)),
183            );
184
185            resample_into(
186                &self.denoised_48k_buf,
187                DENOISE_RATE,
188                TARGET_RATE,
189                &mut self.downsampled_buf,
190            );
191            self.output_16k.extend_from_slice(&self.downsampled_buf);
192        }
193
194        &self.output_16k
195    }
196}
197
198pub struct AudioRecorder {
199    stream: Option<cpal::Stream>,
200    shared: Arc<SharedCaptureState>,
201    current_path: Option<PathBuf>,
202    /// Runtime toggle for noise suppression.
203    noise_suppression: Arc<AtomicBool>,
204    /// Shared denoiser state (created once, reused across callbacks).
205    denoiser: Arc<Mutex<Denoiser>>,
206}
207
208/// Mix multi-channel audio to mono by averaging channels.
209pub fn mix_to_mono(data: &[f32], channels: u32) -> Vec<f32> {
210    if channels <= 1 {
211        return data.to_vec();
212    }
213    data.chunks(channels as usize)
214        .map(|frame| frame.iter().sum::<f32>() / channels as f32)
215        .collect()
216}
217
218/// Convert f32 sample to 16-bit PCM i16, clamping to [-1.0, 1.0].
219/// Uses the standard 32768.0 scale factor for symmetric dynamic range.
220pub fn f32_to_i16(sample: f32) -> i16 {
221    let scaled = (sample.clamp(-1.0, 1.0) * 32768.0) as i32;
222    scaled.clamp(i16::MIN as i32, i16::MAX as i32) as i16
223}
224
225/// The WAV spec Whisper expects: 16kHz, 16-bit, mono PCM.
226pub const WHISPER_WAV_SPEC: WavSpec = WavSpec {
227    channels: 1,
228    sample_rate: 16_000,
229    bits_per_sample: 16,
230    sample_format: SampleFormat::Int,
231};
232
233/// Target sample rate for Whisper input.
234pub const TARGET_RATE: u32 = 16_000;
235
236impl Default for AudioRecorder {
237    fn default() -> Self {
238        Self::new()
239    }
240}
241
242impl AudioRecorder {
243    pub fn new() -> Self {
244        Self {
245            stream: None,
246            shared: Arc::new(SharedCaptureState::new()),
247            current_path: None,
248            noise_suppression: Arc::new(AtomicBool::new(true)),
249            denoiser: Arc::new(Mutex::new(Denoiser::new())),
250        }
251    }
252
253    /// Create a recorder with an explicit noise suppression setting.
254    pub fn with_noise_suppression(enabled: bool) -> Self {
255        Self {
256            noise_suppression: Arc::new(AtomicBool::new(enabled)),
257            ..Self::new()
258        }
259    }
260
261    /// Update the noise suppression toggle at runtime.
262    pub fn set_noise_suppression(&self, enabled: bool) {
263        self.noise_suppression.store(enabled, Ordering::Relaxed);
264    }
265
266    /// Open the microphone and start capturing into the pre-roll buffer.
267    ///
268    /// Call once at app startup so recording starts instantly on hotkey press.
269    /// If the stream is already running this is a no-op.
270    pub fn warm(&mut self) -> Result<()> {
271        if self.stream.is_some() {
272            return Ok(());
273        }
274
275        // Platform hook: on macOS, nudge Bluetooth devices into HFP mode
276        // so the microphone is active when we open the stream.
277        super::activate::prepare_default_input();
278
279        let host = cpal::default_host();
280        let device = host.default_input_device().context("No microphone found")?;
281        self.open_device(device)
282    }
283
284    /// Build and start an input stream on the given device.
285    fn open_device(&mut self, device: cpal::Device) -> Result<()> {
286        let device_name = device
287            .description()
288            .map(|d| d.name().to_string())
289            .unwrap_or_else(|_| "<unknown>".into());
290
291        let supported_config = device
292            .default_input_config()
293            .context("Failed to get default input config")?;
294
295        let native_rate = supported_config.sample_rate();
296        let native_channels = supported_config.channels() as u32;
297
298        log::info!(
299            "Opening audio device: \"{device_name}\" ({native_rate}Hz, {native_channels}ch, {:?})",
300            supported_config.sample_format(),
301        );
302
303        let shared = Arc::clone(&self.shared);
304        let ns_flag = Arc::clone(&self.noise_suppression);
305        let denoiser = Arc::clone(&self.denoiser);
306
307        let stream = device
308            .build_input_stream(
309                &supported_config.into(),
310                move |data: &[f32], _: &cpal::InputCallbackInfo| {
311                    let mono = mix_to_mono(data, native_channels);
312                    let resampled = resample(&mono, native_rate, TARGET_RATE);
313
314                    // Apply noise suppression if enabled
315                    let samples: &[f32] = if ns_flag.load(Ordering::Relaxed) {
316                        if let Ok(mut d) = denoiser.try_lock() {
317                            let denoised = d.process(&resampled);
318                            // SAFETY: denoised borrows d which we hold;
319                            // we only use it within this scope while the lock is held.
320                            // Copy out to avoid holding the lock across buffer writes.
321                            let owned: Vec<f32> = denoised.to_vec();
322                            drop(d);
323                            shared.dispatch_samples(&owned);
324                            return;
325                        }
326                        // Denoiser lock contention — fall through to raw samples
327                        &resampled
328                    } else {
329                        &resampled
330                    };
331
332                    shared.dispatch_samples(samples);
333                },
334                |err| {
335                    log::error!("Audio stream error: {err}");
336                },
337                None,
338            )
339            .context("Failed to build input stream")?;
340
341        stream.play().context("Failed to start audio stream")?;
342        self.stream = Some(stream);
343        log::info!("Microphone warmed up (pre-roll: {PRE_ROLL_MS}ms)");
344
345        Ok(())
346    }
347
348    /// Close the current stream and re-open on the current default input device.
349    fn rewarm(&mut self) -> Result<()> {
350        log::info!("Re-opening audio stream on current default device");
351        self.stream = None;
352        if let Ok(mut ring) = self.shared.pre_roll.lock() {
353            ring.clear();
354        }
355        super::activate::prepare_default_input();
356        let host = cpal::default_host();
357        let device = host.default_input_device().context("No microphone found")?;
358        self.open_device(device)
359    }
360
361    /// Ensure the stream is warm, warming it up if needed.
362    /// If the stream exists but is no longer producing audio (e.g. the
363    /// Bluetooth device disconnected), close and re-open it.
364    fn ensure_warm(&mut self) -> Result<()> {
365        if self.stream.is_some() {
366            // Direct probe: snapshot the counter, wait briefly, check again.
367            // This avoids false positives from stale counters that were set
368            // during a previous recording session.
369            let before = self.shared.callback_count.load(Ordering::Relaxed);
370            std::thread::sleep(std::time::Duration::from_millis(50));
371            let after = self.shared.callback_count.load(Ordering::Relaxed);
372            if after == before {
373                log::warn!("Audio stream appears dead (no callbacks in 50ms), re-opening");
374                self.rewarm()?;
375            }
376        } else {
377            self.warm()?;
378        }
379        Ok(())
380    }
381
382    pub fn start_in_memory(&mut self) -> Result<()> {
383        self.ensure_warm()?;
384
385        self.shared.dropped_samples.store(0, Ordering::Relaxed);
386        if let Ok(mut d) = self.denoiser.lock() {
387            d.reset();
388        }
389
390        // Hold pre_roll lock across the transition to prevent audio samples
391        // from going into pre_roll between the drain and recording=true.
392        if let Ok(mut ring) = self.shared.pre_roll.lock() {
393            if let Ok(mut samples) = self.shared.samples.lock() {
394                samples.clear();
395                samples.extend(ring.drain(..));
396            }
397            // Set recording while still holding pre_roll lock so the audio
398            // callback immediately writes to samples instead of pre_roll.
399            self.shared.recording.store(true, Ordering::Release);
400        }
401
402        self.current_path = None;
403
404        Ok(())
405    }
406
407    /// Stop recording and return the captured samples.
408    /// For in-memory recordings (no WAV file).
409    pub fn stop_samples(&mut self) -> Option<Vec<f32>> {
410        self.shared.recording.store(false, Ordering::Release);
411
412        let dropped = self.shared.dropped_samples.load(Ordering::Relaxed);
413        if dropped > 0 {
414            log::warn!("Dropped {dropped} audio samples due to lock contention during recording");
415        }
416
417        let samples = self.shared.samples.lock().ok().map(|b| b.clone());
418        self.current_path.take();
419        samples.filter(|s| !s.is_empty())
420    }
421
422    pub fn start(&mut self, output_path: &Path) -> Result<()> {
423        self.ensure_warm()?;
424
425        let writer = WavWriter::create(output_path, WHISPER_WAV_SPEC)
426            .context("Failed to create WAV file")?;
427
428        self.current_path = Some(output_path.to_path_buf());
429        self.shared.dropped_samples.store(0, Ordering::Relaxed);
430        if let Ok(mut d) = self.denoiser.lock() {
431            d.reset();
432        }
433
434        // Install the WAV writer before draining pre-roll
435        if let Ok(mut guard) = self.shared.writer.lock() {
436            *guard = Some(writer);
437        }
438
439        // Hold pre_roll lock across the drain and recording flag transition
440        // to prevent audio samples from going into pre_roll during the gap.
441        if let Ok(mut ring) = self.shared.pre_roll.lock() {
442            if let Ok(mut samples) = self.shared.samples.lock() {
443                samples.clear();
444                let pre_roll_data: Vec<f32> = ring.drain(..).collect();
445                samples.extend_from_slice(&pre_roll_data);
446
447                if let Ok(mut guard) = self.shared.writer.lock() {
448                    if let Some(ref mut w) = *guard {
449                        for &sample in &pre_roll_data {
450                            let _ = w.write_sample(f32_to_i16(sample));
451                        }
452                    }
453                }
454            }
455            // Set recording while still holding pre_roll lock so the audio
456            // callback immediately writes to samples instead of pre_roll.
457            self.shared.recording.store(true, Ordering::Release);
458        }
459
460        Ok(())
461    }
462
463    /// Return a copy of samples captured since `start()`, beginning at `offset`.
464    /// Samples are 16 kHz mono f32 in the range \[−1, 1\].
465    #[allow(dead_code)]
466    pub fn snapshot(&self, offset: usize) -> Vec<f32> {
467        if let Ok(buf) = self.shared.samples.lock() {
468            if offset < buf.len() {
469                buf[offset..].to_vec()
470            } else {
471                Vec::new()
472            }
473        } else {
474            Vec::new()
475        }
476    }
477
478    /// Number of 16 kHz samples captured since `start()`.
479    #[allow(dead_code)]
480    pub fn sample_count(&self) -> usize {
481        self.shared.samples.lock().map(|b| b.len()).unwrap_or(0)
482    }
483
484    /// A shared handle to the sample buffer for streaming access.
485    pub fn sample_buffer(&self) -> Arc<Mutex<Vec<f32>>> {
486        Arc::clone(&self.shared.samples)
487    }
488
489    pub fn stop(&mut self) -> Option<PathBuf> {
490        // Transition back to standby (stream stays alive for next recording)
491        self.shared.recording.store(false, Ordering::Release);
492
493        let dropped = self.shared.dropped_samples.load(Ordering::Relaxed);
494        if dropped > 0 {
495            log::warn!("Dropped {dropped} audio samples due to lock contention during recording");
496        }
497
498        // Finalize the WAV file
499        if let Ok(mut guard) = self.shared.writer.lock() {
500            if let Some(writer) = guard.take() {
501                let _ = writer.finalize();
502            }
503        }
504
505        self.current_path.take()
506    }
507}
508
509/// Simple linear interpolation resampler.
510/// Good enough for speech; use `rubato` crate for higher quality if needed.
511pub(crate) fn resample(input: &[f32], from_rate: u32, to_rate: u32) -> Vec<f32> {
512    if from_rate == to_rate {
513        return input.to_vec();
514    }
515
516    let ratio = from_rate as f64 / to_rate as f64;
517    let output_len = (input.len() as f64 / ratio) as usize;
518    let mut output = Vec::with_capacity(output_len);
519
520    for i in 0..output_len {
521        let src_idx = i as f64 * ratio;
522        let idx = src_idx as usize;
523        let frac = src_idx - idx as f64;
524
525        let sample = if idx + 1 < input.len() {
526            input[idx] as f64 * (1.0 - frac) + input[idx + 1] as f64 * frac
527        } else if idx < input.len() {
528            input[idx] as f64
529        } else {
530            0.0
531        };
532
533        output.push(sample as f32);
534    }
535
536    output
537}
538
539/// Allocation-free variant of [`resample`] that writes into a caller-supplied buffer.
540fn resample_into(input: &[f32], from_rate: u32, to_rate: u32, output: &mut Vec<f32>) {
541    output.clear();
542    if from_rate == to_rate {
543        output.extend_from_slice(input);
544        return;
545    }
546
547    let ratio = from_rate as f64 / to_rate as f64;
548    let output_len = (input.len() as f64 / ratio) as usize;
549
550    for i in 0..output_len {
551        let src_idx = i as f64 * ratio;
552        let idx = src_idx as usize;
553        let frac = src_idx - idx as f64;
554
555        let sample = if idx + 1 < input.len() {
556            input[idx] as f64 * (1.0 - frac) + input[idx + 1] as f64 * frac
557        } else if idx < input.len() {
558            input[idx] as f64
559        } else {
560            0.0
561        };
562
563        output.push(sample as f32);
564    }
565}
566
567#[cfg(test)]
568mod tests {
569    use super::*;
570
571    #[test]
572    fn test_resample_same_rate() {
573        let input = vec![0.0, 0.5, 1.0, 0.5, 0.0];
574        let output = resample(&input, 16000, 16000);
575        assert_eq!(output, input);
576    }
577
578    #[test]
579    fn test_resample_downsample() {
580        let input: Vec<f32> = (0..48000).map(|i| i as f32 / 48000.0).collect();
581        let output = resample(&input, 48000, 16000);
582        assert!((output.len() as i64 - 16000).abs() <= 1);
583    }
584
585    #[test]
586    fn test_resample_empty() {
587        let output = resample(&[], 48000, 16000);
588        assert!(output.is_empty());
589    }
590
591    #[test]
592    fn test_resample_upsample() {
593        let input: Vec<f32> = (0..8000).map(|i| i as f32 / 8000.0).collect();
594        let output = resample(&input, 8000, 16000);
595        // Should produce roughly 16000 samples from 8000
596        assert!((output.len() as i64 - 16000).abs() <= 1);
597    }
598
599    #[test]
600    fn test_resample_interpolates() {
601        let input = vec![0.0, 1.0];
602        let output = resample(&input, 2, 4);
603        // Middle values should be interpolated between 0.0 and 1.0
604        assert!(output.len() >= 3);
605        assert!(output[1] > 0.0 && output[1] < 1.0);
606    }
607
608    #[test]
609    fn test_resample_single_sample() {
610        let input = vec![0.5];
611        let output = resample(&input, 16000, 16000);
612        assert_eq!(output, vec![0.5]);
613    }
614
615    #[test]
616    fn test_new_recorder() {
617        let recorder = AudioRecorder::new();
618        assert!(recorder.stream.is_none());
619        assert!(recorder.current_path.is_none());
620    }
621
622    #[test]
623    fn test_stop_without_start() {
624        let mut recorder = AudioRecorder::new();
625        let path = recorder.stop();
626        assert!(path.is_none());
627    }
628
629    // -- mix_to_mono --
630
631    #[test]
632    fn test_mix_to_mono_single_channel() {
633        let data = vec![0.1, 0.2, 0.3];
634        let mono = mix_to_mono(&data, 1);
635        assert_eq!(mono, data);
636    }
637
638    #[test]
639    fn test_mix_to_mono_stereo() {
640        let data = vec![0.0, 1.0, 0.5, 0.5, 1.0, 0.0];
641        let mono = mix_to_mono(&data, 2);
642        assert_eq!(mono.len(), 3);
643        assert!((mono[0] - 0.5).abs() < 0.001);
644        assert!((mono[1] - 0.5).abs() < 0.001);
645        assert!((mono[2] - 0.5).abs() < 0.001);
646    }
647
648    #[test]
649    fn test_mix_to_mono_quad() {
650        let data = vec![1.0, 0.0, 0.0, 0.0]; // 4 channels, 1 frame
651        let mono = mix_to_mono(&data, 4);
652        assert_eq!(mono.len(), 1);
653        assert!((mono[0] - 0.25).abs() < 0.001);
654    }
655
656    #[test]
657    fn test_mix_to_mono_empty() {
658        let mono = mix_to_mono(&[], 2);
659        assert!(mono.is_empty());
660    }
661
662    // -- f32_to_i16 --
663
664    #[test]
665    fn test_f32_to_i16_zero() {
666        assert_eq!(f32_to_i16(0.0), 0);
667    }
668
669    #[test]
670    fn test_f32_to_i16_max() {
671        assert_eq!(f32_to_i16(1.0), 32767);
672    }
673
674    #[test]
675    fn test_f32_to_i16_min() {
676        assert_eq!(f32_to_i16(-1.0), -32768);
677    }
678
679    #[test]
680    fn test_f32_to_i16_clamps_over() {
681        assert_eq!(f32_to_i16(2.0), 32767);
682    }
683
684    #[test]
685    fn test_f32_to_i16_clamps_under() {
686        assert_eq!(f32_to_i16(-2.0), -32768);
687    }
688
689    #[test]
690    fn test_f32_to_i16_half() {
691        let v = f32_to_i16(0.5);
692        assert!(v > 16000 && v < 17000);
693    }
694
695    // -- WHISPER_WAV_SPEC --
696
697    #[test]
698    fn test_whisper_wav_spec() {
699        assert_eq!(WHISPER_WAV_SPEC.channels, 1);
700        assert_eq!(WHISPER_WAV_SPEC.sample_rate, 16_000);
701        assert_eq!(WHISPER_WAV_SPEC.bits_per_sample, 16);
702        assert_eq!(WHISPER_WAV_SPEC.sample_format, SampleFormat::Int);
703    }
704
705    #[test]
706    fn test_target_rate() {
707        assert_eq!(TARGET_RATE, 16_000);
708    }
709
710    #[test]
711    fn test_resample_preserves_endpoints() {
712        let input = vec![0.0, 0.25, 0.5, 0.75, 1.0];
713        let output = resample(&input, 44100, 16000);
714        // First sample should be close to 0.0
715        assert!((output[0] - 0.0).abs() < 0.01);
716    }
717
718    #[test]
719    fn test_resample_large_ratio() {
720        let input: Vec<f32> = (0..96000).map(|i| (i as f32).sin()).collect();
721        let output = resample(&input, 96000, 16000);
722        assert!((output.len() as i64 - 16000).abs() <= 1);
723    }
724
725    // -- snapshot / sample_count --
726
727    #[test]
728    fn test_snapshot_empty_recorder() {
729        let recorder = AudioRecorder::new();
730        assert!(recorder.snapshot(0).is_empty());
731        assert_eq!(recorder.sample_count(), 0);
732    }
733
734    #[test]
735    fn test_sample_buffer_returns_clone() {
736        let recorder = AudioRecorder::new();
737        let buf = recorder.sample_buffer();
738        assert_eq!(buf.lock().unwrap().len(), 0);
739    }
740
741    #[test]
742    fn test_stop_samples_without_start_returns_none() {
743        let mut recorder = AudioRecorder::new();
744        let samples = recorder.stop_samples();
745        assert!(samples.is_none());
746    }
747
748    #[test]
749    fn test_default_trait() {
750        let recorder = AudioRecorder::default();
751        assert!(recorder.stream.is_none());
752        assert!(recorder.current_path.is_none());
753        assert_eq!(recorder.sample_count(), 0);
754    }
755
756    #[test]
757    fn test_snapshot_with_offset_beyond_len() {
758        let recorder = AudioRecorder::new();
759        let snap = recorder.snapshot(100);
760        assert!(snap.is_empty());
761    }
762
763    #[test]
764    fn test_snapshot_with_manual_samples() {
765        let recorder = AudioRecorder::new();
766        // Push samples via the public sample_buffer() handle
767        {
768            let buf = recorder.sample_buffer();
769            buf.lock()
770                .unwrap()
771                .extend_from_slice(&[0.1, 0.2, 0.3, 0.4, 0.5]);
772        }
773        let snap = recorder.snapshot(0);
774        assert_eq!(snap.len(), 5);
775        assert!((snap[0] - 0.1).abs() < 0.001);
776
777        let snap = recorder.snapshot(3);
778        assert_eq!(snap.len(), 2);
779        assert!((snap[0] - 0.4).abs() < 0.001);
780    }
781
782    #[test]
783    fn test_sample_count_with_manual_samples() {
784        let recorder = AudioRecorder::new();
785        assert_eq!(recorder.sample_count(), 0);
786        {
787            let buf = recorder.sample_buffer();
788            buf.lock().unwrap().extend_from_slice(&[0.0; 100]);
789        }
790        assert_eq!(recorder.sample_count(), 100);
791    }
792
793    #[test]
794    fn test_mix_to_mono_six_channels() {
795        // 6-channel surround: one frame
796        let data = vec![1.0, 0.0, 0.0, 0.0, 0.0, 0.0];
797        let mono = mix_to_mono(&data, 6);
798        assert_eq!(mono.len(), 1);
799        assert!((mono[0] - 1.0 / 6.0).abs() < 0.001);
800    }
801
802    #[test]
803    fn test_resample_ratio_accuracy() {
804        // 44.1kHz -> 16kHz: common real-world ratio
805        let input: Vec<f32> = (0..44100).map(|i| (i as f32 / 44100.0).sin()).collect();
806        let output = resample(&input, 44100, 16000);
807        // Should produce approximately 16000 samples
808        assert!((output.len() as i64 - 16000).abs() <= 1);
809    }
810
811    #[test]
812    fn test_f32_to_i16_negative_half() {
813        let v = f32_to_i16(-0.5);
814        assert!(v < -16000 && v > -17000);
815    }
816
817    #[test]
818    fn test_stop_clears_current_path() {
819        let mut recorder = AudioRecorder::new();
820        recorder.current_path = Some(std::path::PathBuf::from("/tmp/test.wav"));
821        let path = recorder.stop();
822        // stop() should return and clear the path
823        assert_eq!(path, Some(std::path::PathBuf::from("/tmp/test.wav")));
824        assert!(recorder.current_path.is_none());
825    }
826
827    // -- pre-roll --
828
829    #[test]
830    fn test_pre_roll_constants() {
831        assert_eq!(PRE_ROLL_MS, 200);
832        assert_eq!(PRE_ROLL_SAMPLES, 3200);
833    }
834
835    #[test]
836    fn test_warm_is_idempotent_without_device() {
837        // warm() will fail without a real audio device, but calling new() is fine
838        let recorder = AudioRecorder::new();
839        assert!(recorder.stream.is_none());
840    }
841
842    // ── Denoiser ──
843
844    #[test]
845    fn test_denoiser_new() {
846        let d = Denoiser::new();
847        assert!(d.pending_16k.is_empty());
848        assert!(d.output_16k.is_empty());
849        assert!(d.first_frame);
850    }
851
852    #[test]
853    fn test_denoiser_reset() {
854        let mut d = Denoiser::new();
855        d.pending_16k.push(1.0);
856        d.output_16k.push(2.0);
857        d.first_frame = false;
858        d.reset();
859        assert!(d.pending_16k.is_empty());
860        assert!(d.output_16k.is_empty());
861        assert!(d.first_frame);
862    }
863
864    #[test]
865    fn test_denoiser_process_empty() {
866        let mut d = Denoiser::new();
867        let out = d.process(&[]);
868        assert!(out.is_empty());
869    }
870
871    #[test]
872    fn test_denoiser_process_short_accumulates() {
873        let mut d = Denoiser::new();
874        let out = d.process(&[0.0; 100]);
875        assert!(out.is_empty());
876        assert_eq!(d.pending_16k.len(), 100);
877    }
878
879    #[test]
880    fn test_denoiser_process_one_frame_skipped() {
881        let mut d = Denoiser::new();
882        // One frame = 160 samples at 16 kHz, but the first frame is always skipped.
883        let out = d.process(&[0.0; 160]);
884        assert!(out.is_empty());
885        assert!(!d.first_frame);
886    }
887
888    #[test]
889    fn test_denoiser_process_two_frames_produces_output() {
890        let mut d = Denoiser::new();
891        // First frame skipped, second frame produces 160 samples.
892        let out = d.process(&[0.0; 320]);
893        assert_eq!(out.len(), 160);
894    }
895
896    #[test]
897    fn test_denoiser_process_multiple_frames() {
898        let mut d = Denoiser::new();
899        // 3 frames: first skipped, remaining 2 produce 320 samples.
900        let out = d.process(&[0.0; 480]);
901        assert_eq!(out.len(), 320);
902    }
903
904    #[test]
905    fn test_denoiser_continuity_across_calls() {
906        let mut d = Denoiser::new();
907        // 100 samples: too few for a frame
908        let out1 = d.process(&[0.1; 100]);
909        assert!(out1.is_empty());
910
911        // 100 more → 200 total, one frame processed (160) but skipped, 40 leftover
912        let out2 = d.process(&[0.1; 100]);
913        assert!(out2.is_empty());
914        assert_eq!(d.pending_16k.len(), 40);
915
916        // 120 more → 40 + 120 = 160, second frame produces output
917        let out3 = d.process(&[0.1; 120]).to_vec();
918        assert_eq!(out3.len(), 160);
919    }
920
921    // ── dispatch_samples ──
922
923    #[test]
924    fn test_dispatch_recording_appends_to_samples() {
925        let state = SharedCaptureState::new();
926        state.recording.store(true, Ordering::Release);
927        state.dispatch_samples(&[0.1, 0.2, 0.3]);
928        let buf = state.samples.lock().unwrap();
929        assert_eq!(&*buf, &[0.1, 0.2, 0.3]);
930    }
931
932    #[test]
933    fn test_dispatch_standby_appends_to_pre_roll() {
934        let state = SharedCaptureState::new();
935        state.dispatch_samples(&[0.5, 0.6]);
936        let ring = state.pre_roll.lock().unwrap();
937        assert_eq!(ring.len(), 2);
938        assert!((ring[0] - 0.5).abs() < f32::EPSILON);
939        assert!((ring[1] - 0.6).abs() < f32::EPSILON);
940    }
941
942    #[test]
943    fn test_dispatch_pre_roll_caps_at_limit() {
944        let state = SharedCaptureState::new();
945        let filler: Vec<f32> = (0..PRE_ROLL_SAMPLES).map(|i| i as f32).collect();
946        state.dispatch_samples(&filler);
947        // Push 2 more — oldest samples should be evicted
948        state.dispatch_samples(&[99.0, 100.0]);
949        let ring = state.pre_roll.lock().unwrap();
950        assert_eq!(ring.len(), PRE_ROLL_SAMPLES);
951        assert!((ring[ring.len() - 1] - 100.0).abs() < f32::EPSILON);
952        assert!((ring[ring.len() - 2] - 99.0).abs() < f32::EPSILON);
953    }
954
955    #[test]
956    fn test_dispatch_with_writer() {
957        let dir = tempfile::tempdir().unwrap();
958        let path = dir.path().join("test.wav");
959        let writer = WavWriter::create(&path, WHISPER_WAV_SPEC).unwrap();
960        let state = SharedCaptureState::new();
961        state.recording.store(true, Ordering::Release);
962        *state.writer.lock().unwrap() = Some(writer);
963
964        state.dispatch_samples(&[0.1, 0.2, 0.3]);
965
966        // Finalize the WAV and verify it was written
967        if let Some(w) = state.writer.lock().unwrap().take() {
968            w.finalize().unwrap();
969        }
970        let reader = hound::WavReader::open(&path).unwrap();
971        let written: Vec<i16> = reader.into_samples::<i16>().map(|s| s.unwrap()).collect();
972        assert_eq!(written.len(), 3);
973    }
974
975    // ── Pre-roll buffer logic ──
976
977    #[test]
978    fn test_pre_roll_does_not_exceed_capacity() {
979        let state = SharedCaptureState::new();
980        let large: Vec<f32> = (0..(PRE_ROLL_SAMPLES + 500)).map(|i| i as f32).collect();
981        state.dispatch_samples(&large);
982        let ring = state.pre_roll.lock().unwrap();
983        assert_eq!(ring.len(), PRE_ROLL_SAMPLES);
984    }
985
986    #[test]
987    fn test_pre_roll_drains_into_samples_on_start() {
988        let recorder = AudioRecorder::new();
989        // Manually push samples into pre_roll
990        {
991            let mut ring = recorder.shared.pre_roll.lock().unwrap();
992            ring.extend([0.1, 0.2, 0.3]);
993        }
994        // Simulate what start_in_memory does (without ensure_warm)
995        {
996            let mut ring = recorder.shared.pre_roll.lock().unwrap();
997            let mut samples = recorder.shared.samples.lock().unwrap();
998            samples.clear();
999            samples.extend(ring.drain(..));
1000            recorder.shared.recording.store(true, Ordering::Release);
1001        }
1002        assert!(recorder.shared.recording.load(Ordering::Acquire));
1003        let buf = recorder.shared.samples.lock().unwrap();
1004        assert_eq!(&*buf, &[0.1, 0.2, 0.3]);
1005        let ring = recorder.shared.pre_roll.lock().unwrap();
1006        assert!(ring.is_empty());
1007    }
1008
1009    // ── Recording state transitions ──
1010
1011    #[test]
1012    fn test_new_recorder_not_recording() {
1013        let recorder = AudioRecorder::new();
1014        assert!(!recorder.shared.recording.load(Ordering::Acquire));
1015    }
1016
1017    #[test]
1018    fn test_start_sets_recording_flag() {
1019        let recorder = AudioRecorder::new();
1020        recorder.shared.recording.store(true, Ordering::Release);
1021        assert!(recorder.shared.recording.load(Ordering::Acquire));
1022    }
1023
1024    #[test]
1025    fn test_stop_samples_returns_samples_and_clears_flag() {
1026        let mut recorder = AudioRecorder::new();
1027        recorder.shared.recording.store(true, Ordering::Release);
1028        recorder
1029            .shared
1030            .samples
1031            .lock()
1032            .unwrap()
1033            .extend_from_slice(&[0.1, 0.2]);
1034        let samples = recorder.stop_samples();
1035        assert!(!recorder.shared.recording.load(Ordering::Acquire));
1036        assert_eq!(samples.unwrap(), vec![0.1, 0.2]);
1037    }
1038
1039    #[test]
1040    fn test_stop_samples_returns_none_when_empty() {
1041        let mut recorder = AudioRecorder::new();
1042        // Not recording and no samples → None
1043        assert!(recorder.stop_samples().is_none());
1044    }
1045
1046    #[test]
1047    fn test_sample_count_zero_initially() {
1048        let recorder = AudioRecorder::new();
1049        assert_eq!(recorder.sample_count(), 0);
1050    }
1051
1052    #[test]
1053    fn test_snapshot_empty_initially() {
1054        let recorder = AudioRecorder::new();
1055        assert!(recorder.snapshot(0).is_empty());
1056    }
1057}