onedrop_engine/

audio_input.rs

//! Real-time audio input capture using cpal.
//!
//! On Linux, cpal sits on top of ALSA and transparently sees both
//! PulseAudio and PipeWire (via its Pulse compatibility layer)
//! devices. The interesting bit for a music visualizer is selecting
//! the right *source*: the host's default input is usually the
//! microphone, but we want the active sink's monitor so the
//! visualizer reacts to whatever is playing out of the speakers.
//! [`detect_default_source`] applies that heuristic — pick the
//! monitor of the default output sink when one exists, otherwise
//! any `.monitor` source, otherwise the host's default input.
//! [`with_device`] uses it transparently when no explicit name is
//! requested.

use cpal::traits::{DeviceTrait, HostTrait, StreamTrait};
use cpal::{Device, Host, Stream, StreamConfig};
use std::sync::{Arc, Mutex};
use thiserror::Error;

/// Audio input errors.
#[derive(Debug, Error)]
pub enum AudioInputError {
    #[error("No audio input device available")]
    NoDevice,

    #[error("Failed to get default input config: {0}")]
    ConfigError(#[from] cpal::DefaultStreamConfigError),

    #[error("Failed to build audio stream: {0}")]
    BuildStreamError(#[from] cpal::BuildStreamError),

    #[error("Failed to play audio stream: {0}")]
    PlayStreamError(#[from] cpal::PlayStreamError),
}

pub type Result<T> = std::result::Result<T, AudioInputError>;

/// Coarse classification of an input source. Used by the CLI
/// `list-audio` subcommand and the GUI device picker to flag which
/// devices are monitor sources (capture what speakers play) versus
/// microphones — both surface as cpal `Device`s with no built-in
/// type tag, so we infer from the name.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SourceKind {
    /// PulseAudio / PipeWire sink monitor (captures system audio).
    Monitor,
    /// Physical or virtual microphone.
    Microphone,
    /// Anything we can't confidently classify — bluetooth headsets,
    /// loopback sinks, virtual mixers, …
    Other,
}

/// A capture device enumerated from the active cpal host.
#[derive(Debug, Clone)]
pub struct AudioSource {
    /// Device name as cpal reports it. Pass back to
    /// [`AudioInput::with_device`] to open this device specifically.
    pub name: String,
    /// Classification heuristic. See [`SourceKind`].
    pub kind: SourceKind,
    /// `true` when this source matches the host's default *input*
    /// device — what `with_device(None)` would historically pick.
    pub is_host_default: bool,
    /// `true` when this source is the monitor of the host's default
    /// *output* sink — what the autodetect logic prefers when no
    /// name is requested. Mutually compatible with `is_host_default`
    /// (rare, but possible if a user wires the monitor as their
    /// default input).
    pub is_autodetect_pick: bool,
}

/// Read a cpal `Device`'s human name through the
/// [`DeviceTrait::description`] surface (cpal 0.17 deprecates the bare
/// `name()` accessor). Returns `None` when the description can't be
/// fetched — usually because the device was disconnected mid-enumeration.
fn device_name(d: &Device) -> Option<String> {
    d.description().ok().map(|desc| desc.name().to_string())
}

/// Classify a device name with the heuristic documented above.
/// Pure function so we can unit-test without a real audio host.
pub(crate) fn classify_source(name: &str) -> SourceKind {
    let lower = name.to_ascii_lowercase();
    // `.monitor` suffix is the canonical Pulse / PipeWire signal —
    // every sink exposes one as `<sink_name>.monitor`.
    if lower.ends_with(".monitor") || lower.contains("monitor of ") {
        return SourceKind::Monitor;
    }
    // ALSA-style microphone names. `alsa_input.` is the Pulse prefix
    // for the underlying ALSA capture device.
    if lower.starts_with("alsa_input.")
        || lower.contains("microphone")
        || lower.contains(" mic")
        || lower.ends_with(" mic")
        || lower == "mic"
    {
        return SourceKind::Microphone;
    }
    SourceKind::Other
}

/// Pick the best autodetect source for a music visualizer.
///
/// Order of preference, mirroring what a user expects after
/// double-clicking the binary while music is already playing:
///
/// 1. Monitor of the host's default *output* sink (so the
///    visualizer reacts to whatever the user is listening to).
/// 2. Any other `.monitor` source (covers headphones-only or
///    USB-DAC setups where the "default" sink is something
///    unexpected).
/// 3. Host default input (typically the microphone) — preserves
///    the historical behaviour when no monitor source is available
///    (e.g. ALSA without `module-loopback`).
fn detect_default_source(host: &Host) -> Option<Device> {
    // Try the monitor of the default output first.
    let default_output_name = host.default_output_device().as_ref().and_then(device_name);
    if let Some(name) = default_output_name.as_deref() {
        let want = format!("{}.monitor", name);
        if let Ok(mut iter) = host.input_devices()
            && let Some(d) = iter.find(|d| device_name(d).as_deref() == Some(want.as_str()))
        {
            log::info!("Audio autodetect: monitor of default sink ({want})");
            return Some(d);
        }
    }

    // Otherwise any monitor source.
    if let Ok(mut iter) = host.input_devices()
        && let Some(d) = iter.find(|d| {
            device_name(d)
                .map(|n| classify_source(&n) == SourceKind::Monitor)
                .unwrap_or(false)
        })
    {
        log::info!(
            "Audio autodetect: first available monitor ({})",
            device_name(&d).unwrap_or_else(|| "Unknown".to_string())
        );
        return Some(d);
    }

    // Finally fall back to the host default input.
    host.default_input_device()
}

/// Enumerate every capture source the active cpal host exposes, with
/// classification flags. Cheap (no streams opened); the CLI / GUI
/// device pickers and the autodetect-preview both call this.
pub fn list_sources() -> Vec<AudioSource> {
    let host = cpal::default_host();
    let host_default = host.default_input_device().as_ref().and_then(device_name);
    let default_output = host.default_output_device().as_ref().and_then(device_name);
    let autodetect_monitor = default_output.as_deref().map(|n| format!("{}.monitor", n));

    let Ok(devices) = host.input_devices() else {
        return Vec::new();
    };
    devices
        .filter_map(|d| {
            let name = device_name(&d)?;
            let kind = classify_source(&name);
            let is_host_default = host_default.as_deref() == Some(name.as_str());
            let is_autodetect_pick =
                autodetect_monitor.as_deref() == Some(name.as_str()) && kind == SourceKind::Monitor;
            Some(AudioSource {
                name,
                kind,
                is_host_default,
                is_autodetect_pick,
            })
        })
        .collect()
}

/// Real-time audio input capture.
pub struct AudioInput {
    /// Audio host
    _host: Host,

    /// Input device
    _device: Device,

    /// Input stream
    _stream: Stream,

    /// Shared buffer for audio samples
    buffer: Arc<Mutex<Vec<f32>>>,

    /// Sample rate
    sample_rate: u32,
}

impl AudioInput {
    /// Enumerate the names of every input device offered by cpal's
    /// default host. Used by the GUI's Options panel to populate the
    /// audio device picker. Devices that fail `name()` (rare —
    /// usually disconnected mid-enumeration) are silently skipped.
    ///
    /// Prefer [`list_sources`] when callers want the
    /// monitor/microphone classification flags (CLI list-audio
    /// subcommand, future GUI picker).
    pub fn list_input_devices() -> Vec<String> {
        list_sources().into_iter().map(|s| s.name).collect()
    }

    /// Create a new audio input capture.
    pub fn new() -> Result<Self> {
        Self::with_device(None)
    }

    /// Create a new audio input capture targeting a specific device by
    /// name. `name = None` triggers the music-visualizer autodetect:
    /// the monitor of the host's default output sink is preferred over
    /// the host default input (the microphone), so a fresh install
    /// reacts to playing music without any GUI / CLI configuration.
    /// See [`detect_default_source`] for the full preference order.
    /// Unknown names log a warning and fall back to the same
    /// autodetect.
    pub fn with_device(name: Option<&str>) -> Result<Self> {
        let host = cpal::default_host();

        // Pick the named device when one is requested, otherwise run
        // the autodetect. Unknown names log a warning and fall back —
        // matches what happens when the user unplugs a saved device.
        let device = if let Some(want) = name {
            let mut matched: Option<Device> = None;
            if let Ok(devices) = host.input_devices() {
                for d in devices {
                    if device_name(&d).as_deref() == Some(want) {
                        matched = Some(d);
                        break;
                    }
                }
            }
            match matched {
                Some(d) => d,
                None => {
                    log::warn!("Input device {want:?} not found, using autodetect");
                    detect_default_source(&host).ok_or(AudioInputError::NoDevice)?
                }
            }
        } else {
            detect_default_source(&host).ok_or(AudioInputError::NoDevice)?
        };

        log::info!(
            "Using audio input device: {}",
            device
                .description()
                .map(|d| d.to_string())
                .unwrap_or_else(|_| "Unknown".to_string())
        );

        // Get default input config
        let config = device.default_input_config()?;
        let sample_rate: u32 = config.sample_rate();

        log::info!(
            "Audio input config: {} Hz, {} channels",
            sample_rate,
            config.channels()
        );

        // Create shared buffer
        let buffer = Arc::new(Mutex::new(Vec::new()));
        let buffer_clone = buffer.clone();

        // Build input stream
        let stream_config: StreamConfig = config.into();
        let stream = device.build_input_stream(
            &stream_config,
            move |data: &[f32], _: &cpal::InputCallbackInfo| {
                // Copy audio data to buffer (handle mutex poisoning gracefully)
                if let Ok(mut buf) = buffer_clone.lock() {
                    buf.clear();
                    buf.extend_from_slice(data);
                }
            },
            |err| {
                log::error!("Audio input stream error: {}", err);
            },
            None,
        )?;

        // Start the stream
        stream.play()?;

        log::info!("Audio input stream started");

        Ok(Self {
            _host: host,
            _device: device,
            _stream: stream,
            buffer,
            sample_rate,
        })
    }

    /// Get the latest audio samples.
    /// Returns a copy of the current buffer.
    pub fn get_samples(&self) -> Vec<f32> {
        self.buffer
            .lock()
            .map(|buf| buf.clone())
            .unwrap_or_default()
    }

    /// Get the sample rate.
    pub fn sample_rate(&self) -> u32 {
        self.sample_rate
    }

    /// Get a fixed number of samples for processing.
    /// If not enough samples are available, returns zeros.
    pub fn get_fixed_samples(&self, count: usize) -> Vec<f32> {
        self.buffer
            .lock()
            .map(|buf| {
                if buf.len() >= count {
                    buf[..count].to_vec()
                } else {
                    // Pad with zeros if not enough samples
                    let mut result = buf.clone();
                    result.resize(count, 0.0);
                    result
                }
            })
            .unwrap_or_else(|_| vec![0.0; count])
    }
}

/// Audio input with FFT analysis for bass/mid/treb extraction.
pub struct AudioAnalysisInput {
    /// Audio input
    input: AudioInput,

    /// FFT planner
    fft: Arc<dyn rustfft::Fft<f32>>,

    /// FFT buffer size
    fft_size: usize,
}

impl AudioAnalysisInput {
    /// Create a new audio analysis input.
    ///
    /// # Arguments
    /// * `fft_size` - Size of FFT window (power of 2, e.g., 2048)
    pub fn new(fft_size: usize) -> Result<Self> {
        let input = AudioInput::new()?;

        // Create FFT planner
        let mut planner = rustfft::FftPlanner::new();
        let fft = planner.plan_fft_forward(fft_size);

        log::info!("Audio analysis initialized with FFT size {}", fft_size);

        Ok(Self {
            input,
            fft,
            fft_size,
        })
    }

    /// Analyze audio and extract bass, mid, treb levels.
    /// Returns (bass, mid, treb) in range [0.0, 1.0].
    pub fn analyze(&self) -> (f32, f32, f32) {
        use rustfft::num_complex::Complex;

        // Get samples
        let samples = self.input.get_fixed_samples(self.fft_size);

        // Convert to complex
        let mut buffer: Vec<Complex<f32>> = samples.iter().map(|&s| Complex::new(s, 0.0)).collect();

        // Apply Hann window
        for (i, sample) in buffer.iter_mut().enumerate() {
            let window =
                0.5 * (1.0 - (2.0 * std::f32::consts::PI * i as f32 / self.fft_size as f32).cos());
            *sample *= window;
        }

        // Perform FFT
        self.fft.process(&mut buffer);

        // Calculate magnitude spectrum
        let magnitudes: Vec<f32> = buffer
            .iter()
            .take(self.fft_size / 2) // Only use first half (Nyquist)
            .map(|c| c.norm())
            .collect();

        // Extract bass, mid, treb
        // Frequency bins: bin_freq = sample_rate * bin_index / fft_size
        let sample_rate = self.input.sample_rate() as f32;
        let _bin_to_freq = |bin: usize| sample_rate * bin as f32 / self.fft_size as f32;

        // Bass: 20-250 Hz
        // Mid: 250-2000 Hz
        // Treb: 2000-20000 Hz
        let bass_end = (250.0 * self.fft_size as f32 / sample_rate) as usize;
        let mid_end = (2000.0 * self.fft_size as f32 / sample_rate) as usize;

        // Bounds checking to prevent panics
        let bass_end = bass_end.max(1).min(magnitudes.len());
        let mid_end = mid_end.max(bass_end).min(magnitudes.len());

        let bass: f32 = if bass_end > 1 {
            magnitudes[1..bass_end].iter().sum::<f32>() / (bass_end - 1) as f32
        } else {
            0.0
        };
        let mid: f32 = if mid_end > bass_end {
            magnitudes[bass_end..mid_end].iter().sum::<f32>() / (mid_end - bass_end) as f32
        } else {
            0.0
        };
        let treb: f32 = if magnitudes.len() > mid_end {
            magnitudes[mid_end..].iter().sum::<f32>() / (magnitudes.len() - mid_end) as f32
        } else {
            0.0
        };

        // Normalize to [0, 1] range (approximate)
        let normalize = |x: f32| (x * 10.0).min(1.0);

        (normalize(bass), normalize(mid), normalize(treb))
    }

    /// Get the sample rate.
    pub fn sample_rate(&self) -> u32 {
        self.input.sample_rate()
    }

    /// Get raw samples for further processing.
    pub fn get_samples(&self) -> Vec<f32> {
        self.input.get_samples()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn classify_monitor_suffix() {
        assert_eq!(
            classify_source("alsa_output.pci-0000_00_1f.3.analog-stereo.monitor"),
            SourceKind::Monitor
        );
        assert_eq!(
            classify_source("Monitor of Built-in Audio"),
            SourceKind::Monitor
        );
        assert_eq!(
            classify_source("hdmi-stereo-extra1.monitor"),
            SourceKind::Monitor
        );
    }

    #[test]
    fn classify_microphone() {
        assert_eq!(
            classify_source("alsa_input.pci-0000_00_1f.3.analog-stereo"),
            SourceKind::Microphone
        );
        assert_eq!(
            classify_source("Internal Microphone"),
            SourceKind::Microphone
        );
        assert_eq!(classify_source("Headset Mic"), SourceKind::Microphone);
    }

    #[test]
    fn classify_other_for_unknown_names() {
        // PipeWire native node names aren't always tagged — we
        // conservatively classify those as Other so the autodetect
        // never *guesses* the wrong source.
        assert_eq!(
            classify_source("Family 17h HD Audio Controller"),
            SourceKind::Other
        );
        assert_eq!(classify_source("USB Audio Device"), SourceKind::Other);
        assert_eq!(classify_source(""), SourceKind::Other);
    }

    #[test]
    fn classify_is_case_insensitive() {
        assert_eq!(classify_source("Foo.MONITOR"), SourceKind::Monitor);
        assert_eq!(
            classify_source("BUILT-IN MICROPHONE"),
            SourceKind::Microphone
        );
    }

    #[test]
    #[ignore] // Requires audio device
    fn test_audio_input_creation() {
        let input = AudioInput::new();
        assert!(input.is_ok(), "Failed to create audio input");
    }

    #[test]
    #[ignore] // Requires audio device
    fn test_audio_analysis_creation() {
        let input = AudioAnalysisInput::new(2048);
        assert!(input.is_ok(), "Failed to create audio analysis input");
    }

    #[test]
    #[ignore] // Requires audio device
    fn test_audio_capture() {
        let input = AudioInput::new().unwrap();

        // Wait a bit for samples
        std::thread::sleep(std::time::Duration::from_millis(100));

        let samples = input.get_samples();
        assert!(!samples.is_empty(), "Should have captured some samples");
    }

    #[test]
    #[ignore] // Requires audio device
    fn test_audio_analysis() {
        let input = AudioAnalysisInput::new(2048).unwrap();

        // Wait a bit for samples
        std::thread::sleep(std::time::Duration::from_millis(100));

        let (bass, mid, treb) = input.analyze();

        // Should be in valid range
        assert!((0.0..=1.0).contains(&bass), "Bass out of range: {}", bass);
        assert!((0.0..=1.0).contains(&mid), "Mid out of range: {}", mid);
        assert!((0.0..=1.0).contains(&treb), "Treb out of range: {}", treb);

        println!(
            "Audio levels - Bass: {:.3}, Mid: {:.3}, Treb: {:.3}",
            bass, mid, treb
        );
    }
}