path: root/src/recording/mod.rs

//! Voice channel recording: a rolling in-memory PCM ring buffer holding the
//! last `window` seconds of mixed audio, flushable to an mp3 on demand.
//!
//! ## How audio arrives
//! songbird delivers a `CoreEvent::VoiceTick` every 20 ms carrying
//! already-decoded, jitter-buffered mono i16 PCM per speaking SSRC — but only
//! when the driver is configured with `DecodeMode::Decode(Channels::Mono,
//! SampleRate::Hz48000)` (set in `main.rs` via `register_songbird_from_config`).
//! That tick is *both* the audio source and the wall-clock driver: there is no
//! separate tokio timer, no `CancellationToken`, no per-SSRC Opus decoder, no
//! f32→s16 conversion. songbird pays the decode cost; we only mix.
//!
//! ## Lifecycle / memory
//! The `EventHandler` registered on the `Call` holds a `Weak<Recorder>`. The
//! sole strong `Arc<Recorder>` lives in `VoiceManager.recorder`. Tearing that
//! slot down (`None`) drops the ring immediately — ~86 MB at 15 min / 48 kHz —
//! even while the songbird `Call` itself lingers. While the slot is `None`,
//! `Weak::upgrade` in the handler returns `None` and the tick is a no-op.
//!
//! ## Ring semantics
//! `head` points at the first sample of the 20 ms window currently being
//! written. Each tick: zero `[head, head+960)` (clearing audio aged out by one
//! window-length, which is how silence is placed when nobody speaks), saturating
//! add every speaker's 960 mono i16 samples into that window, then advance
//! `head += 960` (mod `window_samples`). The newest live sample is therefore
//! `head - 1`; the oldest is `head` (exactly one window old once the ring has
//! wrapped). `flush` reads `min(elapsed, window, written)` samples oldest-→newest.

use std::process::Stdio;
use std::sync::{Arc, Mutex, Weak};

use poise::serenity_prelude as serenity;
use songbird::{Event, EventContext, EventHandler};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tracing::warn;

use crate::error::{Error, Result};

/// The songbird mixer clock fires one `VoiceTick` per 20 ms ⇒ 50 ticks/sec.
const TICKS_PER_SEC: usize = 50;

pub struct Recorder {
    inner: Mutex<Inner>,
    sample_rate: u32,
    samples_per_frame: usize,
    window_samples: usize,
}

struct Inner {
    ring: Vec<i16>,
    /// First sample index of the 20 ms window currently being written.
    head: usize,
    /// Monotonic count of 20 ms frames written since recording started.
    frames_written: u64,
}

impl Recorder {
    pub fn new(sample_rate: u32, window_seconds: u64) -> Self {
        let samples_per_frame = (sample_rate as usize) / TICKS_PER_SEC;
        // Defensive: callers validate config, but a zero-length ring would
        // panic on `% len` in `mix_tick`.
        let window_seconds = window_seconds.max(1);
        let window_samples = (sample_rate as usize) * (window_seconds as usize);
        // `head` advances by `samples_per_frame` and wraps mod `window_samples`,
        // so the window length must be a whole number of frames.
        debug_assert!(
            window_samples % samples_per_frame == 0,
            "window must be a whole number of 20 ms frames"
        );
        Self {
            inner: Mutex::new(Inner {
                ring: vec![0i16; window_samples],
                head: 0,
                frames_written: 0,
            }),
            sample_rate,
            samples_per_frame,
            window_samples,
        }
    }

    pub fn sample_rate(&self) -> u32 {
        self.sample_rate
    }

    /// Snapshot the last `min(frames_written, window)` frames of audio,
    /// oldest→newest, as little-endian mono i16 bytes (the `s16le` format
    /// ffmpeg consumes).
    ///
    /// Read-only: the ring is left intact (flush-and-keep, per the `/record`
    /// design — successive calls may overlap).
    ///
    /// `frames_written` is authoritative: the ring holds exactly the last
    /// `min(frames_written, TICKS_PER_SEC * window_seconds)` frames, each
    /// written by one real 20 ms tick. A wall-clock term would drift from the
    /// mixer whenever ticks stall or bunch; the ring is the ground truth.
    pub fn flush(&self) -> Vec<u8> {
        let inner = match self.inner.lock() {
            Ok(g) => g,
            Err(e) => {
                warn!("recorder mutex poisoned: {e}");
                return Vec::new();
            }
        };
        let len = self.window_samples;
        let d = inner
            .frames_written
            .saturating_mul(self.samples_per_frame as u64)
            .min(len as u64) as usize;
        if d == 0 {
            return Vec::new();
        }
        // Oldest live sample sits at `head - d` (mod len); read `d` forward.
        let start = (inner.head + len - d) % len;
        let mut bytes = Vec::with_capacity(d * 2);
        if start + d <= len {
            push_le_i16(&mut bytes, &inner.ring[start..start + d]);
        } else {
            let first = len - start;
            push_le_i16(&mut bytes, &inner.ring[start..]);
            push_le_i16(&mut bytes, &inner.ring[..d - first]);
        }
        bytes
    }

    /// Absorb one 20 ms tick: zero the window about to be written (silence
    /// placement), saturating-add each speaker's decoded mono i16 frame, then
    /// advance `head` by one frame.
    ///
    /// Split out from the `EventHandler` so tests can drive the ring without
    /// constructing songbird's `#[non_exhaustive]` `VoiceTick`.
    pub fn mix_tick(&self, speakers: &[Vec<i16>]) {
        let mut inner = match self.inner.lock() {
            Ok(g) => g,
            Err(_) => return,
        };
        let spf = self.samples_per_frame;
        let len = self.window_samples;

        // Zero the window about to be written. This is the silence-placement
        // step: it clears samples that aged out one window-length ago, so a
        // tick with no speakers writes pure silence rather than stale audio.
        for i in 0..spf {
            let pos = (inner.head + i) % len;
            inner.ring[pos] = 0;
        }
        // Mix every speaking SSRC's decoded mono i16 into the window.
        // `saturating_add` prevents wraparound when several speakers overlap
        // (the only real failure mode on the hot path); multi-speaker overlaps
        // may clip slightly, the correct trade for a 50 Hz × N-packet mix.
        for audio in speakers {
            let n = audio.len().min(spf);
            for i in 0..n {
                let pos = (inner.head + i) % len;
                inner.ring[pos] = inner.ring[pos].saturating_add(audio[i]);
            }
        }
        inner.head = (inner.head + spf) % len;
        inner.frames_written += 1;
    }
}

fn push_le_i16(out: &mut Vec<u8>, samples: &[i16]) {
    for s in samples {
        out.extend_from_slice(&s.to_le_bytes());
    }
}

/// `songbird::EventHandler` registered on each `Call` while recording is
/// enabled. Holds only a `Weak<Recorder>`, so dropping the strong `Arc` in
/// `VoiceManager.recorder` immediately frees the ring; subsequent ticks no-op.
pub struct RecorderEventHandler {
    weak: Weak<Recorder>,
}

impl RecorderEventHandler {
    pub fn new(recorder: &Arc<Recorder>) -> Self {
        Self {
            weak: Arc::downgrade(recorder),
        }
    }
}
#[async_trait::async_trait]
impl EventHandler for RecorderEventHandler {
    async fn act(&self, ctx: &EventContext<'_>) -> Option<Event> {
        // Only registered for `CoreEvent::VoiceTick`, but match defensively.
        let tick = match ctx {
            EventContext::VoiceTick(t) => t,
            _ => return None,
        };
        let Some(rec) = self.weak.upgrade() else {
            // Recorder torn down — stop doing work, stay registered (cheap no-op)
            // until songbird drops the Call and us along with it.
            return None;
        };
        let frames: Vec<Vec<i16>> = tick
            .speaking
            .values()
            .filter_map(|vd| vd.decoded_voice.clone())
            .collect();
        rec.mix_tick(&frames);
        None // keep listening
    }
}

/// Encode little-endian mono s16 PCM to mp3 via an ffmpeg subprocess.
///
/// Mirrors the bot's other ffmpeg-subprocess patterns (e.g. `ffmpeg_stream_input`)
/// but for a *finite* buffer: write all PCM to stdin (EOF on drop), drain stderr
/// so ffmpeg can't block on a full pipe, collect mp3 from stdout.
pub async fn encode_mp3(pcm_s16le: &[u8], sample_rate: u32) -> Result<Vec<u8>> {
    let mut child = tokio::process::Command::new("ffmpeg")
        .args([
            "-y",
            "-f",
            "s16le",
            "-ar",
            &sample_rate.to_string(),
            "-ac",
            "1",
            "-i",
            "pipe:0",
            // Discord playback levels land quiet in recordings; 2x loudness.
            // Output-side option, so it sits after `-i`. Clipping on loud
            // overlaps is acceptable for voice.
            "-af",
            "volume=2.0",
            "-c:a",
            "libmp3lame",
            "-b:a",
            "64k",
            "-f",
            "mp3",
            "pipe:1",
        ])
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .kill_on_drop(true)
        .spawn()
        .map_err(|e| Error::Custom(format!("ffmpeg spawn: {e}")))?;

    let stdin = child.stdin.take();
    let stdout = child
        .stdout
        .take()
        .ok_or_else(|| Error::Custom("ffmpeg stdout not captured".into()))?;
    let stderr = child.stderr.take();

    let pcm_owned = pcm_s16le.to_vec();
    let write_task = tokio::spawn(async move {
        if let Some(mut stdin) = stdin {
            let _ = stdin.write_all(&pcm_owned).await;
            // `stdin` drops here ⇒ ffmpeg sees EOF on its input.
        }
    });

    // Drain stderr in the background so a full pipe can't deadlock ffmpeg,
    // and surface the trace in the logs if encoding fails.
    if let Some(stderr) = stderr {
        tokio::spawn(async move {
            let mut buf = Vec::with_capacity(4096);
            let mut reader = tokio::io::BufReader::new(stderr);
            let _ = reader.read_to_end(&mut buf).await;
            if !buf.is_empty() {
                warn!(
                    "ffmpeg (mp3 encode) stderr: {}",
                    String::from_utf8_lossy(&buf).trim_end()
                );
            }
        });
    }

    let mut mp3 = Vec::new();
    let mut reader = tokio::io::BufReader::new(stdout);
    reader
        .read_to_end(&mut mp3)
        .await
        .map_err(|e| Error::Custom(format!("ffmpeg read: {e}")))?;

    let _ = write_task.await;
    let status = child
        .wait()
        .await
        .map_err(|e| Error::Custom(format!("ffmpeg wait: {e}")))?;
    if !status.success() {
        return Err(Error::Custom(
            "ffmpeg mp3 encode failed (non-zero exit)".into(),
        ));
    }
    Ok(mp3)
}

/// Reconcile the recorder with the bot's own voice state.
///
/// Called from `events::handle_voice_state_update` for the bot itself (before
/// the existing bot-skip). The explicit-join path (`VoiceManager::connect`)
/// starts the recorder itself; this handles the cases `connect` can't see:
/// admin moves (re-setup if the recorder was dropped) and all leaves —
/// `/stop`, kicks, gateway drops — which tear it down.
pub async fn on_bot_voice_state(
    voice_manager: &crate::voice::VoiceManager,
    guild_id: Option<serenity::GuildId>,
    old_channel: Option<serenity::ChannelId>,
    new_channel: Option<serenity::ChannelId>,
) {
    if !voice_manager.recording_enabled() {
        return;
    }
    let Some(guild_id) = guild_id else {
        return;
    };
    // Already torn down by `VoiceManager::disconnect` in the `/stop` path; this
    // is the catch-all for leaves `disconnect` didn't initiate (kicks, drops).
    if new_channel.is_none() {
        voice_manager.stop_recording(guild_id).await;
        return;
    }
    // Joined or moved into a channel. If this guild's recorder is missing
    // (e.g. a move broke the prior Call), try to (re)register on the current
    // Call.
    if new_channel != old_channel && voice_manager.get_recorder(guild_id).await.is_none() {
        if let Some(call) = voice_manager.get_call(guild_id) {
            voice_manager.start_recording(guild_id, &call).await;
        }
    }
}

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

    /// Decode the `s16le` byte stream `flush` produces back into i16 samples.
    fn to_i16(bytes: &[u8]) -> Vec<i16> {
        bytes
            .chunks_exact(2)
            .map(|c| i16::from_le_bytes([c[0], c[1]]))
            .collect()
    }

    #[test]
    fn empty_recorder_flushes_nothing() {
        let rec = Recorder::new(48_000, 1);
        assert!(rec.flush().is_empty());
    }

    #[test]
    fn zero_window_is_clamped_and_does_not_panic() {
        let rec = Recorder::new(48_000, 0);
        let spf = rec.samples_per_frame;
        rec.mix_tick(&[vec![1i16; spf]]);
        assert!(!rec.flush().is_empty());
    }

    #[test]
    fn single_speaker_values_survive() {
        let rec = Recorder::new(48_000, 1);
        let spf = rec.samples_per_frame;
        // 10 ticks of a constant 100.
        for _ in 0..10 {
            rec.mix_tick(&[vec![100i16; spf]]);
        }
        let out = to_i16(&rec.flush());
        assert_eq!(out.len(), 10 * spf);
        assert!(out.iter().all(|&s| s == 100));
    }

    #[test]
    fn silent_tick_writes_silence_not_stale_audio() {
        let rec = Recorder::new(48_000, 1);
        let spf = rec.samples_per_frame;
        // Fill the whole 50-frame window with 100s.
        for _ in 0..TICKS_PER_SEC {
            rec.mix_tick(&[vec![100i16; spf]]);
        }
        // A tick with no speakers must age out one frame of the 100s…
        rec.mix_tick(&[]);
        let out = to_i16(&rec.flush());
        assert_eq!(out.len(), TICKS_PER_SEC * spf);
        // …and write 960 pure-silence samples at the *newest* end, keeping
        // 49 frames of the old audio as the oldest content.
        assert!(out[out.len() - spf..].iter().all(|&s| s == 0));
        assert!(out[..(TICKS_PER_SEC - 1) * spf].iter().all(|&s| s == 100));
    }

    #[test]
    fn two_speakers_mix_saturating() {
        let rec = Recorder::new(48_000, 1);
        let spf = rec.samples_per_frame;
        rec.mix_tick(&[vec![20000i16; spf], vec![10000i16; spf]]);
        let out = to_i16(&rec.flush());
        assert!(out.iter().all(|&s| s == 30000));

        // Overflow saturates instead of wrapping.
        let rec = Recorder::new(48_000, 1);
        rec.mix_tick(&[vec![30000i16; spf], vec![30000i16; spf]]);
        let out = to_i16(&rec.flush());
        assert!(out.iter().all(|&s| s == i16::MAX));
    }

    #[test]
    fn ring_wraps_and_flush_orders_oldest_to_newest() {
        // 2-frame window keeps the wrap tight and the assertion unambiguous.
        // window_samples = 2 * spf ⇒ a whole number of 20 ms frames.
        let rec = Recorder::new(48_000, 0); // 0 s → 0 samples; not a whole frame…
        // …so build one manually with a 2-frame window instead.
        drop(rec);
        let spf = 960usize;
        let rec = Recorder {
            inner: Mutex::new(Inner {
                ring: vec![0i16; 2 * spf],
                head: 0,
                frames_written: 0,
            }),
            sample_rate: 48_000,
            samples_per_frame: spf,
            window_samples: 2 * spf,
        };
        // Frame A = 7s, then frame B = 9s, then frame C = 11s (overwrites A).
        rec.mix_tick(&[vec![7i16; spf]]);
        rec.mix_tick(&[vec![9i16; spf]]);
        rec.mix_tick(&[vec![11i16; spf]]);
        let out = to_i16(&rec.flush());
        assert_eq!(out.len(), 2 * spf);
        // Oldest→newest must be B, C — the wrapped-over A is gone.
        assert!(out[..spf].iter().all(|&s| s == 9));
        assert!(out[spf..].iter().all(|&s| s == 11));
    }

    #[test]
    fn flush_is_read_only() {
        let rec = Recorder::new(48_000, 1);
        let spf = rec.samples_per_frame;
        rec.mix_tick(&[vec![50i16; spf]]);
        let first = to_i16(&rec.flush());
        let second = to_i16(&rec.flush());
        assert_eq!(first, second);
    }

    #[tokio::test]
    async fn encode_mp3_produces_mp3_frames() {
        // 1 s of a 440 Hz sine at 48 kHz mono.
        let sample_rate = 48_000u32;
        let mut pcm = Vec::with_capacity(sample_rate as usize * 2);
        for n in 0..sample_rate {
            let t = n as f64 / sample_rate as f64;
            let s = (2.0 * std::f64::consts::PI * 440.0 * t).sin() * 8000.0;
            pcm.extend_from_slice(&(s as i16).to_le_bytes());
        }
        let mp3 = encode_mp3(&pcm, sample_rate)
            .await
            .expect("mp3 encode failed");
        assert!(!mp3.is_empty(), "ffmpeg produced no mp3 bytes");
        // MP3 frame sync: 0xFF followed by a byte with the upper 3 bits set
        // (0xE0 mask), or an ID3 tag header.
        let has_frame = mp3.windows(2).any(|w| w[0] == 0xFF && (w[1] & 0xE0) == 0xE0);
        let has_id3 = mp3.starts_with(b"ID3");
        assert!(has_frame || has_id3, "output is not an mp3 stream");
    }
}