Lumotia/crates/audio/src/wav.rs

use std::io::BufWriter;
use std::path::Path;

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

/// Append-friendly WAV writer for long-running captures.
///
/// The in-memory `Vec<f32>` used by `run_live_session` to persist audio
/// on session end (brief item #19) has three failure modes: (a) a crash
/// during transcription takes the recording with it; (b) RAM bloat at
/// long session lengths; (c) an OOM kills the capture loop. `WavWriter`
/// replaces that pattern with an on-disk writer that periodically
/// flushes the WAV header so the file on disk is a valid, playable WAV
/// at any point the process is interrupted.
///
/// The writer samples at the rate / channel count supplied at
/// construction; callers read those from
/// `LocalEngine::capabilities()` (brief item #13 wiring) rather than
/// hardcoding 16 kHz / mono.
pub struct WavWriter {
    inner: hound::WavWriter<BufWriter<std::fs::File>>,
    samples_since_flush: usize,
    flush_every: usize,
}

impl WavWriter {
    /// Sample count between automatic header flushes. Flushing costs
    /// two seeks per call; 8000 samples at 16 kHz = 500 ms, so the
    /// worst-case "last half second is lost on crash" bound holds.
    const DEFAULT_FLUSH_EVERY_SAMPLES: usize = 8_000;

    /// Create a new WAV file at `path`, truncating any previous content.
    /// Header reflects zero samples until the first `flush` or
    /// `finalize`.
    pub fn create(path: &Path, sample_rate: u32, channels: u16) -> Result<Self> {
        let spec = hound::WavSpec {
            channels,
            sample_rate,
            bits_per_sample: 16,
            sample_format: hound::SampleFormat::Int,
        };
        let file = std::fs::File::create(path).map_err(KonError::Io)?;
        let buffered = BufWriter::new(file);
        let inner = hound::WavWriter::new(buffered, spec)
            .map_err(|e| KonError::Io(std::io::Error::other(format!("WAV create failed: {e}"))))?;
        Ok(Self {
            inner,
            samples_since_flush: 0,
            flush_every: Self::DEFAULT_FLUSH_EVERY_SAMPLES,
        })
    }

    /// Append f32 samples in `[-1.0, 1.0]`. Samples outside that range
    /// are clamped (matching `write_wav`). Automatically flushes the
    /// header every `flush_every` samples so the on-disk file stays a
    /// valid WAV even if the process is killed between appends.
    pub fn append(&mut self, samples: &[f32]) -> Result<()> {
        for &sample in samples {
            let clamped = sample.clamp(-1.0, 1.0);
            let int_sample = (clamped * i16::MAX as f32) as i16;
            self.inner.write_sample(int_sample).map_err(|e| {
                KonError::Io(std::io::Error::other(format!("WAV write failed: {e}")))
            })?;
        }
        self.samples_since_flush += samples.len();
        if self.samples_since_flush >= self.flush_every {
            self.flush()?;
        }
        Ok(())
    }

    /// Force an immediate header flush. Leaves the file in a valid-WAV
    /// state up to the current sample count. Callers do not need to
    /// call this explicitly — `append` flushes every
    /// `Self::DEFAULT_FLUSH_EVERY_SAMPLES` — but may do so at natural
    /// boundaries (end-of-utterance, UI events) for tighter recovery.
    pub fn flush(&mut self) -> Result<()> {
        self.inner
            .flush()
            .map_err(|e| KonError::Io(std::io::Error::other(format!("WAV flush failed: {e}"))))?;
        self.samples_since_flush = 0;
        Ok(())
    }

    /// Finalise the WAV: writes the terminal header state and closes
    /// the file. Call on clean session end. A dropped-without-finalize
    /// writer leaves a playable file up to the last flush; callers
    /// that care about the unflushed tail should always finalise.
    pub fn finalize(self) -> Result<()> {
        self.inner.finalize().map_err(|e| {
            KonError::Io(std::io::Error::other(format!("WAV finalize failed: {e}")))
        })?;
        Ok(())
    }
}

/// Write f32 PCM samples to a 16-bit WAV file.
pub fn write_wav(path: &Path, audio: &AudioSamples) -> Result<()> {
    let spec = hound::WavSpec {
        channels: audio.channels(),
        sample_rate: audio.sample_rate(),
        bits_per_sample: 16,
        sample_format: hound::SampleFormat::Int,
    };

    let mut writer = hound::WavWriter::create(path, spec)
        .map_err(|e| KonError::Io(std::io::Error::other(format!("WAV create failed: {e}"))))?;

    for &sample in audio.samples() {
        let clamped = sample.clamp(-1.0, 1.0);
        let int_sample = (clamped * i16::MAX as f32) as i16;
        writer
            .write_sample(int_sample)
            .map_err(|e| KonError::Io(std::io::Error::other(format!("WAV write failed: {e}"))))?;
    }

    writer
        .finalize()
        .map_err(|e| KonError::Io(std::io::Error::other(format!("WAV finalize failed: {e}"))))?;

    Ok(())
}

/// Read a WAV file to f32 PCM `AudioSamples`.
///
/// Any per-sample decode error is surfaced as `KonError::AudioDecodeFailed`
/// rather than silently dropped. A previous implementation used
/// `filter_map(|s| s.ok())`, so a truncated or corrupt payload returned
/// a short, silently-partial `AudioSamples` — callers got `Ok` while
/// losing audio (flagged by the 2026-04-22 review).
pub fn read_wav(path: &Path) -> Result<AudioSamples> {
    let reader = hound::WavReader::open(path)
        .map_err(|e| KonError::AudioDecodeFailed(format!("WAV open failed: {e}")))?;

    let spec = reader.spec();
    let sample_rate = spec.sample_rate;
    let channels = spec.channels;
    let bits_per_sample = spec.bits_per_sample;

    let samples: Vec<f32> = match spec.sample_format {
        hound::SampleFormat::Int => reader
            .into_samples::<i32>()
            .map(|sample| {
                sample
                    .map(|s| s as f32 / (1 << (bits_per_sample - 1)) as f32)
                    .map_err(|e| {
                        KonError::AudioDecodeFailed(format!("WAV sample decode failed: {e}"))
                    })
            })
            .collect::<Result<Vec<f32>>>()?,
        hound::SampleFormat::Float => reader
            .into_samples::<f32>()
            .map(|sample| {
                sample.map_err(|e| {
                    KonError::AudioDecodeFailed(format!("WAV sample decode failed: {e}"))
                })
            })
            .collect::<Result<Vec<f32>>>()?,
    };

    Ok(AudioSamples::new(samples, sample_rate, channels))
}

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

    #[test]
    fn wav_writer_survives_crash() {
        // Property under test: a `WavWriter` that has been flushed but
        // never finalised leaves a valid, readable WAV on disk. This
        // is the crash-safety guarantee — if the kon process aborts
        // mid-session, the on-disk file up to the last flush is
        // recoverable.
        //
        // `std::mem::forget` is the canonical way to simulate an
        // abort inside a unit test: it skips the Drop impl (which
        // would otherwise finalise the hound writer for us) and
        // mirrors what happens when the OS reaps the process without
        // giving Rust a chance to run destructors.
        let temp_dir = std::env::temp_dir();
        let path = temp_dir.join("kon_test_wav_writer_survives_crash.wav");
        let _ = std::fs::remove_file(&path);

        let mut writer = WavWriter::create(&path, 16_000, 1).unwrap();
        let flushed_samples = vec![0.1_f32; 16_000]; // 1s
        writer.append(&flushed_samples).unwrap();
        writer.flush().unwrap();

        // Post-flush, append another second that will NOT be reflected
        // in the header if the writer dies before the next flush.
        let unflushed_tail = vec![0.2_f32; 16_000];
        writer.append(&unflushed_tail).unwrap();

        // Abort — Drop does not run, the hound finaliser is skipped.
        std::mem::forget(writer);

        let loaded = read_wav(&path).unwrap();
        assert_eq!(loaded.sample_rate(), 16_000);
        assert!(
            loaded.samples().len() >= 16_000,
            "expected at least the flushed 16000 samples, got {}",
            loaded.samples().len()
        );
        // The flushed portion is readable and approximately correct.
        for s in &loaded.samples()[..16_000] {
            assert!(
                (s - 0.1).abs() < 0.01,
                "flushed sample {s} deviates from 0.1 beyond 16-bit quantisation slack",
            );
        }

        let _ = std::fs::remove_file(&path);
    }

    #[test]
    fn wav_writer_append_then_finalize_roundtrips() {
        let temp_dir = std::env::temp_dir();
        let path = temp_dir.join("kon_test_wav_writer_finalize.wav");
        let _ = std::fs::remove_file(&path);

        let mut writer = WavWriter::create(&path, 16_000, 1).unwrap();
        writer.append(&vec![0.0_f32; 8_000]).unwrap();
        writer.append(&vec![0.5_f32; 8_000]).unwrap();
        writer.finalize().unwrap();

        let loaded = read_wav(&path).unwrap();
        assert_eq!(loaded.sample_rate(), 16_000);
        assert_eq!(loaded.samples().len(), 16_000);

        let _ = std::fs::remove_file(&path);
    }

    #[test]
    fn read_wav_surfaces_truncated_sample_stream_errors() {
        // Regression for the 2026-04-22 review: filter_map(|s| s.ok())
        // previously swallowed decode errors on corrupt input, so a
        // truncated WAV returned Ok with a short samples vec. The
        // new code must propagate the error.
        let temp_dir = std::env::temp_dir();
        let path = temp_dir.join("kon_test_truncated_wav.wav");
        let _ = std::fs::remove_file(&path);

        // Write 100 samples (200 bytes at 16-bit).
        let original = AudioSamples::mono_16khz((0..100).map(|i| (i as f32) / 100.0).collect());
        write_wav(&path, &original).unwrap();

        // Drop the last 10 bytes — 5 samples' worth. hound's iterator
        // should surface an UnexpectedEof on the final read once its
        // internal data-chunk accounting runs out of bytes.
        let content = std::fs::read(&path).unwrap();
        let truncated = &content[..content.len() - 10];
        std::fs::write(&path, truncated).unwrap();

        let result = read_wav(&path);
        assert!(
            result.is_err(),
            "truncated WAV must surface an AudioDecodeFailed error, got: {result:?}"
        );

        let _ = std::fs::remove_file(&path);
    }

    #[test]
    fn wav_roundtrip() {
        let temp_dir = std::env::temp_dir();
        let path = temp_dir.join("kon_test_roundtrip.wav");

        let original = AudioSamples::mono_16khz(vec![0.0, 0.5, -0.5, 0.25, -0.25]);
        write_wav(&path, &original).unwrap();

        let loaded = read_wav(&path).unwrap();
        assert_eq!(loaded.sample_rate(), 16000);
        assert_eq!(loaded.samples().len(), 5);

        // 16-bit quantisation introduces small error
        for (a, b) in original.samples().iter().zip(loaded.samples().iter()) {
            assert!(
                (a - b).abs() < 0.001,
                "Sample mismatch: original={a}, loaded={b}"
            );
        }

        std::fs::remove_file(&path).ok();
    }
}