peerspeak/src/audio/multitrack.rs

//! Multitrack (stem) recording: one synced WAV per peer + your mic, plus an
//! optional convenience mixed track — podcast/streamer-grade source for
//! post-production. See `docs/multitrack-recording-plan.md`.
//!
//! All tracks share one master clock: the playout-mixer cycle. Every cycle,
//! **exactly `FRAME_SAMPLES` samples are appended to every track** — real audio
//! for peers who produced a frame that cycle, silence for those idle — so every
//! track stays sample-aligned by construction. A peer that joins mid-recording
//! has its track pre-padded with silence back to cycle 0, so all stems line up
//! at sample 0 on a timeline.
//!
//! This module is pure plumbing over [`WavWriter`]: no audio decode, no
//! networking, no realtime work. The mixer (a non-RT task) drives it.

use std::collections::{HashMap, VecDeque};
use std::io;
use std::path::{Path, PathBuf};

use iroh::EndpointId;

use crate::audio::recorder::WavWriter;
use crate::core::jitter::FRAME_SAMPLES;

/// Cap on the silence chunk written at once when pre-padding a late joiner, so a
/// long-running call can't trigger a single multi-hundred-MB allocation.
const SILENCE_CHUNK: usize = FRAME_SAMPLES * 256;

/// Cap on buffered mic samples (~200ms @ 48kHz). Bounds how far the mic track
/// can drift if the capture clock runs ahead of the mixer cycle; past it the
/// oldest mic audio is dropped. Mirrors `recorder::MAX_MIC_FIFO`.
const MAX_MIC_FIFO: usize = 48_000 / 5;
const MAX_SESSION_DIR_ATTEMPTS: usize = 1_000;

/// Create a collision-free session directory for a timestamp. The base
/// timestamp is tried first, followed by `-2`, `-3`, and so on; an existing
/// recording is never reopened or overwritten.
pub fn create_session_dir(base: &Path, now_unix_secs: u64) -> io::Result<PathBuf> {
    let filename = crate::audio::recorder::timestamp_filename(now_unix_secs);
    let stem = filename.trim_end_matches(".wav");
    for attempt in 1..=MAX_SESSION_DIR_ATTEMPTS {
        let name = if attempt == 1 {
            stem.to_string()
        } else {
            format!("{stem}-{attempt}")
        };
        let path = base.join(name);
        match std::fs::create_dir(&path) {
            Ok(()) => return Ok(path),
            Err(e) if e.kind() == io::ErrorKind::AlreadyExists => continue,
            Err(e) => return Err(e),
        }
    }
    Err(io::Error::new(
        io::ErrorKind::AlreadyExists,
        "multitrack directory suffixes exhausted",
    ))
}

/// One output track: its WAV writer plus whether it has been written *this*
/// cycle (so `end_cycle` knows which tracks to pad with silence).
struct Track {
    writer: WavWriter,
    written_this_cycle: bool,
}

impl Track {
    fn create(path: &Path) -> io::Result<Self> {
        Ok(Self {
            writer: WavWriter::new(path)?,
            written_this_cycle: false,
        })
    }

    /// Append `frame` fitted to exactly `frame_samples` (zero-padded if short),
    /// and mark the track as written for this cycle.
    fn write_frame(&mut self, frame: &[i16], frame_samples: usize) -> io::Result<()> {
        self.writer.write_samples(&fit(frame, frame_samples))?;
        self.written_this_cycle = true;
        Ok(())
    }

    /// Append `samples` of silence (no cycle-marking — used for padding).
    fn write_silence(&mut self, samples: usize) -> io::Result<()> {
        let mut remaining = samples;
        while remaining > 0 {
            let n = remaining.min(SILENCE_CHUNK);
            self.writer.write_samples(&vec![0i16; n])?;
            remaining -= n;
        }
        Ok(())
    }
}

/// Return `frame` resized to exactly `n` samples: truncated if longer (shouldn't
/// happen — Opus frames are uniform), zero-padded if shorter.
fn fit(frame: &[i16], n: usize) -> Vec<i16> {
    let mut v = Vec::with_capacity(n);
    let take = frame.len().min(n);
    v.extend_from_slice(&frame[..take]);
    v.resize(n, 0);
    v
}

/// A filesystem-safe stem filename: a slug of the (already untrusted-sanitized)
/// display name plus a short id suffix to disambiguate same-named peers, e.g.
/// `alice-3b1f9c2a.wav`. Falls back to `peer` when the name slugs to nothing.
pub fn track_filename(name: &str, id: &EndpointId) -> String {
    let clean = crate::sanitize::sanitize_name(name);
    let mut slug: String = clean
        .chars()
        .map(|c| if c.is_ascii_alphanumeric() { c.to_ascii_lowercase() } else { '-' })
        .collect();
    // Collapse runs of '-' and trim them off the ends.
    while slug.contains("--") {
        slug = slug.replace("--", "-");
    }
    let slug = slug.trim_matches('-');
    let slug = if slug.is_empty() { "peer" } else { slug };
    let short: String = id.to_string().chars().take(8).collect();
    format!("{slug}-{short}.wav")
}

/// A live multitrack recording: per-peer stems + your mic, plus an optional
/// mixed track, all under one session directory and clocked together.
pub struct MultitrackRecorder {
    dir: PathBuf,
    frame_samples: usize,
    /// Cycles recorded so far = the shared length (in frames) of every track.
    cycles: u64,
    peers: HashMap<EndpointId, Track>,
    /// Your mic track. Fed asynchronously from the capture thread via
    /// [`push_mic`](MultitrackRecorder::push_mic) into `mic_fifo`, then drained
    /// one frame per `end_cycle` so it aligns with the cycle clock.
    mic: WavWriter,
    mic_fifo: VecDeque<i16>,
    /// Present in "Both" mode (stems + mixed), absent in "stems only".
    mix: Option<Track>,
}

impl MultitrackRecorder {
    /// Create a recording in `dir` (which must already exist). `with_mix` adds
    /// the convenience mixed track (`mix.wav`). Your mic is always `me.wav`.
    pub fn create(dir: &Path, frame_samples: usize, with_mix: bool) -> io::Result<Self> {
        let mic = WavWriter::new(&dir.join("me.wav"))?;
        let mix = if with_mix {
            Some(Track::create(&dir.join("mix.wav"))?)
        } else {
            None
        };
        Ok(Self {
            dir: dir.to_path_buf(),
            frame_samples,
            cycles: 0,
            peers: HashMap::new(),
            mic,
            mic_fifo: VecDeque::new(),
            mix,
        })
    }

    /// The session directory holding all the track files.
    pub fn dir(&self) -> &Path {
        &self.dir
    }

    /// Register a peer's stem track, pre-padding it with silence back to cycle 0
    /// so it aligns with the others. Idempotent: a peer already tracked is left
    /// as-is (re-announce / name change doesn't restart their file).
    pub fn add_peer(&mut self, id: EndpointId, name: &str) -> io::Result<()> {
        if self.peers.contains_key(&id) {
            return Ok(());
        }
        let mut track = Track::create(&self.dir.join(track_filename(name, &id)))?;
        track.write_silence(self.cycles as usize * self.frame_samples)?;
        self.peers.insert(id, track);
        Ok(())
    }

    /// Record one peer's decoded frame for the current cycle. If the peer wasn't
    /// registered yet (write raced ahead of the join event), auto-register it
    /// with an id-only name so no audio is dropped.
    pub fn write_peer(&mut self, id: EndpointId, frame: &[i16]) -> io::Result<()> {
        if !self.peers.contains_key(&id) {
            self.add_peer(id, "")?;
        }
        let fs = self.frame_samples;
        self.peers.get_mut(&id).unwrap().write_frame(frame, fs)
    }

    /// Buffer a frame of your transmitted mic audio (called from the capture
    /// thread, asynchronously to the mixer cycle). Bounded: oldest samples drop
    /// past [`MAX_MIC_FIFO`] so the mic track's offset can't grow without limit.
    pub fn push_mic(&mut self, frame: &[i16]) {
        self.mic_fifo.extend(frame.iter().copied());
        let overflow = self.mic_fifo.len().saturating_sub(MAX_MIC_FIFO);
        if overflow > 0 {
            self.mic_fifo.drain(..overflow);
        }
    }

    /// Pull exactly `n` mic samples from the FIFO, silence-padded on underrun.
    fn drain_mic(&mut self, n: usize) -> Vec<i16> {
        let take = n.min(self.mic_fifo.len());
        let mut v: Vec<i16> = self.mic_fifo.drain(..take).collect();
        v.resize(n, 0);
        v
    }

    /// Record the finished mixed-bus frame for the current cycle (no-op in
    /// stems-only mode).
    pub fn write_mix(&mut self, frame: &[i16]) -> io::Result<()> {
        let fs = self.frame_samples;
        if let Some(mix) = self.mix.as_mut() {
            mix.write_frame(frame, fs)?;
        }
        Ok(())
    }

    /// Close out the current cycle: every track that wasn't written this cycle
    /// gets one frame of silence, so all tracks advance in lockstep. Call once
    /// per mixer cycle, after the per-track writes.
    pub fn end_cycle(&mut self) -> io::Result<()> {
        let fs = self.frame_samples;
        // Mic: always one frame per cycle, drained from the FIFO (silence on
        // underrun), so it tracks the cycle clock like the peer stems.
        let mic_frame = self.drain_mic(fs);
        self.mic.write_samples(&mic_frame)?;
        // Peers + the optional mix track: pad any not written this cycle.
        for track in self.peers.values_mut().chain(self.mix.as_mut()) {
            if !track.written_this_cycle {
                track.write_silence(fs)?;
            }
            track.written_this_cycle = false;
        }
        self.cycles += 1;
        Ok(())
    }

    /// Finalize every track's WAV header. Consumes the recorder.
    pub fn finalize(self) -> io::Result<()> {
        self.mic.finalize()?;
        if let Some(mix) = self.mix {
            mix.writer.finalize()?;
        }
        for (_, track) in self.peers {
            track.writer.finalize()?;
        }
        Ok(())
    }
}

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

    fn an_id() -> EndpointId {
        SecretKey::generate().public()
    }

    /// Samples of PCM data in a finished WAV file: (len - 44-byte header) / 2.
    fn wav_samples(path: &Path) -> usize {
        let len = std::fs::metadata(path).unwrap().len() as usize;
        (len - 44) / 2
    }

    fn tmpdir(tag: &str) -> PathBuf {
        let d = std::env::temp_dir().join(format!("ps-mt-{}-{}", tag, std::process::id()));
        let _ = std::fs::remove_dir_all(&d);
        std::fs::create_dir_all(&d).unwrap();
        d
    }

    #[test]
    fn fit_pads_and_truncates() {
        assert_eq!(fit(&[1, 2], 4), vec![1, 2, 0, 0]);
        assert_eq!(fit(&[1, 2, 3, 4], 2), vec![1, 2]);
        assert_eq!(fit(&[], 3), vec![0, 0, 0]);
    }

    #[test]
    fn track_filename_is_fs_safe_and_disambiguated() {
        let id = an_id();
        let short: String = id.to_string().chars().take(8).collect();
        assert_eq!(track_filename("Alice", &id), format!("alice-{short}.wav"));
        // Spaces / punctuation collapse to single dashes, trimmed.
        assert_eq!(track_filename("  Bob the Builder! ", &id), format!("bob-the-builder-{short}.wav"));
        // A name that sanitizes/slugs to nothing falls back to "peer".
        assert_eq!(track_filename("!!!", &id), format!("peer-{short}.wav"));
    }

    #[test]
    fn same_second_sessions_get_unique_directories_without_reuse() {
        let base = tmpdir("collision");
        let first = create_session_dir(&base, 1_700_000_000).unwrap();
        std::fs::write(first.join("sentinel"), b"keep me").unwrap();

        let second = create_session_dir(&base, 1_700_000_000).unwrap();

        assert_ne!(second, first);
        assert_eq!(std::fs::read(first.join("sentinel")).unwrap(), b"keep me");
        let _ = std::fs::remove_dir_all(&base);
    }

    #[test]
    fn all_tracks_equal_length_after_n_cycles() {
        let dir = tmpdir("equal");
        let frame = 4; // tiny frame for the test
        let p1 = an_id();
        let p2 = an_id();
        let mut rec = MultitrackRecorder::create(&dir, frame, true).unwrap();
        rec.add_peer(p1, "p1").unwrap();
        rec.add_peer(p2, "p2").unwrap();

        // 3 cycles; p1 talks every cycle, p2 only on cycle 2, mic pushed twice.
        for c in 0..3 {
            rec.write_peer(p1, &[1, 1, 1, 1]).unwrap();
            if c == 2 {
                rec.write_peer(p2, &[2, 2, 2, 2]).unwrap();
            }
            if c < 2 {
                rec.push_mic(&[9, 9, 9, 9]);
            }
            rec.write_mix(&[5, 5, 5, 5]).unwrap();
            rec.end_cycle().unwrap();
        }
        rec.finalize().unwrap();

        let expected = 3 * frame;
        assert_eq!(wav_samples(&dir.join("me.wav")), expected, "mic padded to full length");
        assert_eq!(wav_samples(&dir.join("mix.wav")), expected);
        assert_eq!(wav_samples(&dir.join(track_filename("p1", &p1))), expected);
        assert_eq!(wav_samples(&dir.join(track_filename("p2", &p2))), expected, "silent peer still full length");
    }

    #[test]
    fn late_joiner_is_silence_padded_to_start() {
        let dir = tmpdir("late");
        let frame = 4;
        let early = an_id();
        let late = an_id();
        let mut rec = MultitrackRecorder::create(&dir, frame, false).unwrap();
        rec.add_peer(early, "early").unwrap();

        // 2 cycles before the late peer joins.
        for _ in 0..2 {
            rec.write_peer(early, &[1, 1, 1, 1]).unwrap();
            rec.end_cycle().unwrap();
        }
        // Late peer joins at cycle 2.
        rec.add_peer(late, "late").unwrap();
        for _ in 0..3 {
            rec.write_peer(early, &[1, 1, 1, 1]).unwrap();
            rec.write_peer(late, &[2, 2, 2, 2]).unwrap();
            rec.end_cycle().unwrap();
        }
        rec.finalize().unwrap();

        // Both tracks are the full 5 cycles long (late one was back-padded).
        assert_eq!(wav_samples(&dir.join(track_filename("early", &early))), 5 * frame);
        assert_eq!(wav_samples(&dir.join(track_filename("late", &late))), 5 * frame);

        // The late track's first 2 cycles are silence, then the real audio.
        let bytes = std::fs::read(dir.join(track_filename("late", &late))).unwrap();
        let data = &bytes[44..];
        let read_sample = |i: usize| i16::from_le_bytes([data[i * 2], data[i * 2 + 1]]);
        for i in 0..(2 * frame) {
            assert_eq!(read_sample(i), 0, "leading silence at sample {i}");
        }
        assert_eq!(read_sample(2 * frame), 2, "real audio starts at cycle 2");
    }

    #[test]
    fn stems_only_writes_no_mix_file() {
        let dir = tmpdir("nomix");
        let mut rec = MultitrackRecorder::create(&dir, 4, false).unwrap();
        rec.write_mix(&[1, 2, 3, 4]).unwrap(); // no-op
        rec.end_cycle().unwrap();
        rec.finalize().unwrap();
        assert!(dir.join("me.wav").exists());
        assert!(!dir.join("mix.wav").exists(), "no mix track in stems-only mode");
    }
}