use crate::config::Config;
use std::fs::File;
use std::io::{Read, Seek, SeekFrom};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
const TARGET_RATE: u32 = 16_000;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Encoding {
F32,
I16,
}
impl Encoding {
fn bytes_per_sample(self) -> usize {
match self {
Encoding::F32 => 4,
Encoding::I16 => 2,
}
}
}
#[derive(Debug, Clone, Copy)]
struct StemFormat {
channels: u16,
sample_rate: u32,
encoding: Encoding,
}
#[derive(Debug)]
pub(crate) struct StemTail {
path: PathBuf,
format: StemFormat,
cursor: u64,
carry: Vec<u8>,
resample_pos: f64,
}
impl StemTail {
pub(crate) fn open(path: &Path) -> Result<Self, String> {
let mut file = File::open(path).map_err(|e| format!("open {}: {e}", path.display()))?;
let mut header = [0_u8; 4096];
let read = file
.read(&mut header)
.map_err(|e| format!("read header {}: {e}", path.display()))?;
let (format, data_offset) = parse_header(&header[..read])?;
Ok(Self {
path: path.to_path_buf(),
format,
cursor: data_offset,
carry: Vec::new(),
resample_pos: 0.0,
})
}
pub(crate) fn poll(&mut self) -> Result<Vec<f32>, String> {
let mut file =
File::open(&self.path).map_err(|e| format!("reopen {}: {e}", self.path.display()))?;
let len = file
.metadata()
.map_err(|e| format!("stat {}: {e}", self.path.display()))?
.len();
if len <= self.cursor {
return Ok(Vec::new());
}
let want = (len - self.cursor) as usize;
let mut fresh = vec![0_u8; want];
file.seek(SeekFrom::Start(self.cursor))
.map_err(|e| format!("seek {}: {e}", self.path.display()))?;
let got = file
.read(&mut fresh)
.map_err(|e| format!("read {}: {e}", self.path.display()))?;
fresh.truncate(got);
self.cursor += got as u64;
if !self.carry.is_empty() {
let mut joined = std::mem::take(&mut self.carry);
joined.extend_from_slice(&fresh);
fresh = joined;
}
let frame = self.format.channels as usize * self.format.encoding.bytes_per_sample();
if frame == 0 {
return Err("stem reports zero-width frames".into());
}
let usable = fresh.len() - (fresh.len() % frame);
self.carry = fresh[usable..].to_vec();
Ok(self.decode_to_mono_16k(&fresh[..usable]))
}
fn decode_to_mono_16k(&mut self, bytes: &[u8]) -> Vec<f32> {
let channels = self.format.channels as usize;
let width = self.format.encoding.bytes_per_sample();
let frames = bytes.len() / (channels * width);
let ratio = self.format.sample_rate as f64 / TARGET_RATE as f64;
let mut out = Vec::with_capacity(((frames as f64) / ratio).ceil() as usize + 1);
for frame in 0..frames {
let mut sum = 0.0_f32;
for ch in 0..channels {
let at = (frame * channels + ch) * width;
sum += match self.format.encoding {
Encoding::F32 => {
f32::from_le_bytes([bytes[at], bytes[at + 1], bytes[at + 2], bytes[at + 3]])
}
Encoding::I16 => {
i16::from_le_bytes([bytes[at], bytes[at + 1]]) as f32 / 32768.0
}
};
}
let mono = sum / channels as f32;
if self.resample_pos <= frame as f64 {
out.push(mono);
self.resample_pos += ratio;
}
}
self.resample_pos = (self.resample_pos - frames as f64).max(0.0);
out
}
}
fn parse_header(bytes: &[u8]) -> Result<(StemFormat, u64), String> {
if bytes.len() < 12 || &bytes[0..4] != b"RIFF" || &bytes[8..12] != b"WAVE" {
return Err("not a RIFF/WAVE stem".into());
}
let mut at = 12_usize;
let mut format: Option<StemFormat> = None;
while at + 8 <= bytes.len() {
let id = &bytes[at..at + 4];
let size = u32::from_le_bytes([bytes[at + 4], bytes[at + 5], bytes[at + 6], bytes[at + 7]])
as usize;
let body = at + 8;
if id == b"fmt " {
if body + 16 > bytes.len() {
return Err("fmt chunk truncated".into());
}
let tag = u16::from_le_bytes([bytes[body], bytes[body + 1]]);
let channels = u16::from_le_bytes([bytes[body + 2], bytes[body + 3]]);
let sample_rate = u32::from_le_bytes([
bytes[body + 4],
bytes[body + 5],
bytes[body + 6],
bytes[body + 7],
]);
let bits = u16::from_le_bytes([bytes[body + 14], bytes[body + 15]]);
let encoding = match (tag, bits) {
(3, 32) | (0xFFFE, 32) => Encoding::F32,
(1, 16) | (0xFFFE, 16) => Encoding::I16,
_ => return Err(format!("unsupported stem format (tag {tag}, {bits} bits)")),
};
if channels == 0 || channels > 32 || sample_rate == 0 {
return Err("implausible stem format".into());
}
format = Some(StemFormat {
channels,
sample_rate,
encoding,
});
} else if id == b"data" {
let format = format.ok_or_else(|| "data chunk before fmt".to_string())?;
return Ok((format, body as u64));
}
at = body + size + (size & 1);
}
Err("no data chunk in header window".into())
}
const POLL_INTERVAL: std::time::Duration = std::time::Duration::from_millis(250);
const HEADER_WAIT: std::time::Duration = std::time::Duration::from_secs(30);
pub fn spawn_live_transcription_from_stems(
voice_stem: PathBuf,
system_stem: Option<PathBuf>,
config: &Config,
stop_flag: Arc<AtomicBool>,
) -> Option<std::thread::JoinHandle<()>> {
let (live_tx, sidecar_handle) = crate::capture::start_live_sidecar(config, &stop_flag);
let live_tx = live_tx?;
let feeder = std::thread::Builder::new()
.name("stem-tail-feeder".into())
.spawn(move || {
feed_from_stems(voice_stem, system_stem, &live_tx, &stop_flag);
drop(live_tx);
});
match feeder {
Ok(_) => sidecar_handle,
Err(error) => {
tracing::warn!(%error, "could not start stem tail feeder; live transcript unavailable");
sidecar_handle
}
}
}
fn feed_from_stems(
voice_stem: PathBuf,
system_stem: Option<PathBuf>,
live_tx: &std::sync::mpsc::SyncSender<Vec<f32>>,
stop_flag: &Arc<AtomicBool>,
) {
let Some(mut voice) = wait_for_stem(&voice_stem, stop_flag) else {
return;
};
let mut system = system_stem
.as_deref()
.and_then(|path| wait_for_stem(path, stop_flag));
let mut voice_pending: Vec<f32> = Vec::new();
let mut system_pending: Vec<f32> = Vec::new();
while !stop_flag.load(Ordering::Relaxed) {
match voice.poll() {
Ok(samples) => voice_pending.extend_from_slice(&samples),
Err(error) => tracing::debug!(%error, "voice stem poll failed; continuing"),
}
if let Some(system) = system.as_mut() {
match system.poll() {
Ok(samples) => system_pending.extend_from_slice(&samples),
Err(error) => tracing::debug!(%error, "system stem poll failed; continuing"),
}
}
let chunk = take_mixed(&mut voice_pending, &mut system_pending, system.is_some());
if !chunk.is_empty() && live_tx.send(chunk).is_err() {
return;
}
std::thread::sleep(POLL_INTERVAL);
}
if let Ok(samples) = voice.poll() {
voice_pending.extend_from_slice(&samples);
}
if let Some(system) = system.as_mut() {
if let Ok(samples) = system.poll() {
system_pending.extend_from_slice(&samples);
}
}
let tail = take_mixed(&mut voice_pending, &mut system_pending, system.is_some());
if !tail.is_empty() {
let _ = live_tx.send(tail);
}
}
fn take_mixed(voice: &mut Vec<f32>, system: &mut Vec<f32>, has_system: bool) -> Vec<f32> {
if !has_system {
return std::mem::take(voice);
}
let n = voice.len().min(system.len());
if n == 0 {
return Vec::new();
}
let mixed: Vec<f32> = voice
.drain(..n)
.zip(system.drain(..n))
.map(|(a, b)| (a + b).clamp(-1.0, 1.0))
.collect();
mixed
}
fn wait_for_stem(path: &Path, stop_flag: &Arc<AtomicBool>) -> Option<StemTail> {
let deadline = std::time::Instant::now() + HEADER_WAIT;
while std::time::Instant::now() < deadline {
if stop_flag.load(Ordering::Relaxed) {
return None;
}
match StemTail::open(path) {
Ok(tail) => return Some(tail),
Err(_) => std::thread::sleep(POLL_INTERVAL),
}
}
tracing::warn!(
stem = %path.display(),
"stem header never appeared; live transcript will not run for this capture"
);
None
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Write;
fn write_growing_stem(path: &Path, channels: u16, rate: u32, encoding: Encoding) {
let mut f = File::create(path).unwrap();
let (tag, bits) = match encoding {
Encoding::F32 => (3_u16, 32_u16),
Encoding::I16 => (1_u16, 16_u16),
};
let block_align = channels * (bits / 8);
let byte_rate = rate * block_align as u32;
f.write_all(b"RIFF").unwrap();
f.write_all(&4088_u32.to_le_bytes()).unwrap(); f.write_all(b"WAVE").unwrap();
f.write_all(b"fmt ").unwrap();
f.write_all(&16_u32.to_le_bytes()).unwrap();
f.write_all(&tag.to_le_bytes()).unwrap();
f.write_all(&channels.to_le_bytes()).unwrap();
f.write_all(&rate.to_le_bytes()).unwrap();
f.write_all(&byte_rate.to_le_bytes()).unwrap();
f.write_all(&block_align.to_le_bytes()).unwrap();
f.write_all(&bits.to_le_bytes()).unwrap();
f.write_all(b"data").unwrap();
f.write_all(&0_u32.to_le_bytes()).unwrap(); }
fn append_f32(path: &Path, samples: &[f32]) {
let mut f = File::options().append(true).open(path).unwrap();
for s in samples {
f.write_all(&s.to_le_bytes()).unwrap();
}
}
#[test]
fn reads_only_what_was_appended_since_the_last_poll() {
let dir = tempfile::tempdir().unwrap();
let stem = dir.path().join("voice.wav");
write_growing_stem(&stem, 1, 16_000, Encoding::F32);
let mut tail = StemTail::open(&stem).unwrap();
assert!(tail.poll().unwrap().is_empty(), "nothing written yet");
append_f32(&stem, &[0.1, 0.2, 0.3, 0.4]);
assert_eq!(tail.poll().unwrap().len(), 4);
assert!(tail.poll().unwrap().is_empty());
append_f32(&stem, &[0.5, 0.6]);
assert_eq!(tail.poll().unwrap().len(), 2);
}
#[test]
fn carries_a_partial_frame_across_polls() {
let dir = tempfile::tempdir().unwrap();
let stem = dir.path().join("voice.wav");
write_growing_stem(&stem, 1, 16_000, Encoding::F32);
let mut tail = StemTail::open(&stem).unwrap();
{
let mut f = File::options().append(true).open(&stem).unwrap();
f.write_all(&[1, 2, 3]).unwrap();
}
assert!(
tail.poll().unwrap().is_empty(),
"a partial frame must not be decoded"
);
{
let mut f = File::options().append(true).open(&stem).unwrap();
f.write_all(&[4]).unwrap();
}
assert_eq!(
tail.poll().unwrap().len(),
1,
"the carried bytes complete one frame"
);
}
#[test]
fn downmixes_channels_rather_than_dropping_one() {
let dir = tempfile::tempdir().unwrap();
let stem = dir.path().join("system.wav");
write_growing_stem(&stem, 2, 16_000, Encoding::F32);
let mut tail = StemTail::open(&stem).unwrap();
append_f32(&stem, &[0.0, 1.0]);
let out = tail.poll().unwrap();
assert_eq!(out.len(), 1);
assert!(
(out[0] - 0.5).abs() < 1e-6,
"expected the average, got {out:?}"
);
}
#[test]
fn decimates_to_16k_and_keeps_phase_across_chunks() {
let dir = tempfile::tempdir().unwrap();
let stem = dir.path().join("voice.wav");
write_growing_stem(&stem, 1, 48_000, Encoding::F32);
let mut tail = StemTail::open(&stem).unwrap();
append_f32(&stem, &vec![0.25_f32; 48]);
let first = tail.poll().unwrap().len();
append_f32(&stem, &vec![0.25_f32; 48]);
let second = tail.poll().unwrap().len();
assert!((15..=17).contains(&first), "first chunk: {first}");
assert!((15..=17).contains(&second), "second chunk: {second}");
}
#[test]
fn decodes_i16_stems() {
let dir = tempfile::tempdir().unwrap();
let stem = dir.path().join("voice.wav");
write_growing_stem(&stem, 1, 16_000, Encoding::I16);
let mut tail = StemTail::open(&stem).unwrap();
let mut f = File::options().append(true).open(&stem).unwrap();
f.write_all(&16384_i16.to_le_bytes()).unwrap();
drop(f);
let out = tail.poll().unwrap();
assert_eq!(out.len(), 1);
assert!((out[0] - 0.5).abs() < 1e-3, "got {out:?}");
}
#[test]
fn mixing_waits_for_both_stems_and_buffers_the_surplus() {
let mut voice = vec![0.1, 0.2, 0.3];
let mut system = vec![0.4];
let mixed = take_mixed(&mut voice, &mut system, true);
assert_eq!(mixed.len(), 1, "only the overlapping prefix is emitted");
assert!((mixed[0] - 0.5).abs() < 1e-6);
assert_eq!(voice.len(), 2, "unmatched voice audio stays buffered");
assert!(system.is_empty());
}
#[test]
fn mixing_passes_voice_through_when_there_is_no_system_stem() {
let mut voice = vec![0.1, 0.2];
let mut system = Vec::new();
let out = take_mixed(&mut voice, &mut system, false);
assert_eq!(out, vec![0.1, 0.2]);
assert!(voice.is_empty(), "everything is consumed");
}
#[test]
fn mixing_clamps_instead_of_wrapping() {
let mut voice = vec![0.9];
let mut system = vec![0.9];
let out = take_mixed(&mut voice, &mut system, true);
assert_eq!(out, vec![1.0]);
}
#[test]
fn waiting_for_a_stem_gives_up_when_asked_to_stop() {
let dir = tempfile::tempdir().unwrap();
let missing = dir.path().join("never-created.wav");
let stop = Arc::new(AtomicBool::new(true));
let start = std::time::Instant::now();
let result = wait_for_stem(&missing, &stop);
assert!(result.is_none());
assert!(
start.elapsed() < std::time::Duration::from_secs(2),
"should observe the stop flag immediately, took {:?}",
start.elapsed()
);
}
#[test]
fn rejects_a_header_that_is_not_a_wav() {
let dir = tempfile::tempdir().unwrap();
let stem = dir.path().join("bogus.wav");
std::fs::write(&stem, b"not a wav at all").unwrap();
assert!(StemTail::open(&stem).is_err());
}
}