Lumotia/crates/audio/src/resample.rs

use rubato::{SincFixedIn, SincInterpolationParameters, SincInterpolationType, Resampler, WindowFunction};

use kon_core::constants::WHISPER_SAMPLE_RATE;
use kon_core::error::{KonError, Result};
use kon_core::types::AudioSamples;

/// Resample audio to 16kHz mono using sinc interpolation (rubato).
/// Returns a new AudioSamples at the target sample rate.
pub fn resample_to_16khz(audio: &AudioSamples) -> Result<AudioSamples> {
    let from_rate = audio.sample_rate();
    let target_rate = WHISPER_SAMPLE_RATE;

    if from_rate == target_rate {
        return Ok(AudioSamples::mono_16khz(audio.samples().to_vec()));
    }

    if from_rate == 0 {
        return Err(KonError::AudioDecodeFailed(
            "Cannot resample: source rate is 0".into(),
        ));
    }

    let ratio = target_rate as f64 / from_rate as f64;
    let chunk_size = 1024;

    let params = SincInterpolationParameters {
        sinc_len: 256,
        f_cutoff: 0.95,
        oversampling_factor: 128,
        interpolation: SincInterpolationType::Cubic,
        window: WindowFunction::Blackman,
    };

    let mut resampler = SincFixedIn::<f32>::new(
        ratio,
        1.1,
        params,
        chunk_size,
        1, // mono
    )
    .map_err(|e| {
        KonError::AudioDecodeFailed(format!("Resampler init failed: {e}"))
    })?;

    let samples = audio.samples();
    let mut output_samples: Vec<f32> = Vec::new();

    let mut offset = 0;
    while offset < samples.len() {
        let end = (offset + chunk_size).min(samples.len());
        let mut chunk = samples[offset..end].to_vec();

        if chunk.len() < chunk_size {
            chunk.resize(chunk_size, 0.0);
        }

        let input = vec![chunk];
        let result = resampler.process(&input, None).map_err(|e| {
            KonError::AudioDecodeFailed(format!("Resample failed: {e}"))
        })?;

        if !result.is_empty() && !result[0].is_empty() {
            output_samples.extend_from_slice(&result[0]);
        }

        offset += chunk_size;
    }

    // Trim to expected length (padding may have added extra samples)
    let expected_len = (samples.len() as f64 * ratio) as usize;
    output_samples.truncate(expected_len);

    Ok(AudioSamples::mono_16khz(output_samples))
}

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

    #[test]
    fn resample_passthrough_at_16khz() {
        let input = AudioSamples::mono_16khz(vec![0.1, 0.2, 0.3]);
        let output = resample_to_16khz(&input).unwrap();
        assert_eq!(output.sample_rate(), 16000);
        assert_eq!(output.samples().len(), 3);
    }

    #[test]
    fn resample_preserves_approximate_duration() {
        let rate = 48000;
        let duration_secs = 1.0;
        let num_samples = (rate as f64 * duration_secs) as usize;
        let samples: Vec<f32> =
            (0..num_samples).map(|i| (i as f32 * 0.001).sin()).collect();

        let input = AudioSamples::new(samples, rate, 1);
        let output = resample_to_16khz(&input).unwrap();

        let output_duration = output.samples().len() as f64 / 16000.0;
        assert!(
            (output_duration - duration_secs).abs() < 0.1,
            "Duration mismatch: expected ~{duration_secs}s, got {output_duration}s"
        );
    }
}