peerspeak/src/codec/opus_impl.rs

use crate::codec::{AudioDecoder, AudioEncoder, CodecError};
use crate::config::AudioProfile;
use opus::{Application, Bitrate, Channels, Decoder, Encoder};

/// Concrete libopus encoder settings derived from an [`AudioProfile`]. Plain
/// data, so the profile→params mapping ([`opus_params`]) stays a pure,
/// unit-testable function (W12).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct OpusParams {
    /// Target bitrate in bits/sec.
    pub bitrate: i32,
    /// Enable in-band forward error correction (loss redundancy in the bitstream).
    pub inband_fec: bool,
    /// Expected packet-loss percentage (0..=100); tunes how much FEC libopus adds.
    pub packet_loss_perc: i32,
    /// Discontinuous transmission: stop sending during silence to save bandwidth.
    pub dtx: bool,
}

/// Map a named profile to concrete Opus parameters. Pure — the W12 testable seam.
///
/// `BadNetwork` deliberately runs a *lower* bitrate than `Balanced`: in-band FEC
/// redundancy is carried inside the same bitstream, so trimming the base bitrate
/// leaves headroom for the redundancy on a congested link.
pub fn opus_params(profile: AudioProfile) -> OpusParams {
    match profile {
        AudioProfile::LowLatency => OpusParams {
            bitrate: 24_000,
            inband_fec: false,
            packet_loss_perc: 0,
            dtx: false,
        },
        AudioProfile::Balanced => OpusParams {
            bitrate: 32_000,
            inband_fec: true,
            packet_loss_perc: 10,
            dtx: false,
        },
        AudioProfile::BadNetwork => OpusParams {
            bitrate: 20_000,
            inband_fec: true,
            packet_loss_perc: 25,
            dtx: false,
        },
    }
}

pub struct OpusEncoder {
    encoder: Encoder,
}

impl OpusEncoder {
    /// Creates a new Opus encoder.
    /// Standard voice parameters: sample_rate = 48000, channels = Channels::Mono, application = Application::Voip
    pub fn new(
        sample_rate: u32,
        channels: Channels,
        application: Application,
    ) -> Result<Self, CodecError> {
        let encoder = Encoder::new(sample_rate, channels, application)
            .map_err(|e| CodecError::Init(format!("Failed to create Opus encoder: {}", e)))?;
        Ok(Self { encoder })
    }

    /// Apply concrete codec parameters to the live encoder. Safe to call between
    /// frames, so the user can switch profile mid-call.
    pub fn apply_params(&mut self, params: &OpusParams) -> Result<(), CodecError> {
        self.encoder
            .set_bitrate(Bitrate::Bits(params.bitrate))
            .map_err(|e| CodecError::Init(format!("set_bitrate: {}", e)))?;
        self.encoder
            .set_inband_fec(params.inband_fec)
            .map_err(|e| CodecError::Init(format!("set_inband_fec: {}", e)))?;
        self.encoder
            .set_packet_loss_perc(params.packet_loss_perc)
            .map_err(|e| CodecError::Init(format!("set_packet_loss_perc: {}", e)))?;
        self.encoder
            .set_dtx(params.dtx)
            .map_err(|e| CodecError::Init(format!("set_dtx: {}", e)))?;
        Ok(())
    }

    /// Apply a named [`AudioProfile`] (shorthand for `apply_params(&opus_params(p))`).
    pub fn apply_profile(&mut self, profile: AudioProfile) -> Result<(), CodecError> {
        self.apply_params(&opus_params(profile))
    }
}

impl AudioEncoder for OpusEncoder {
    fn encode(&mut self, pcm: &[i16]) -> Result<Vec<u8>, CodecError> {
        // We allocate a buffer for the compressed output.
        // A maximum packet size of 4000 bytes is more than enough for a single voice frame.
        let mut compressed = vec![0u8; 4000];
        let len = self
            .encoder
            .encode(pcm, &mut compressed)
            .map_err(|e| CodecError::Encode(format!("Opus encoding failed: {}", e)))?;

        compressed.truncate(len);
        Ok(compressed)
    }
}

pub struct OpusDecoder {
    decoder: Decoder,
    channels: Channels,
    /// Samples-per-channel of the frames we transmit (20ms @ 48kHz = 960).
    /// Used to size the Packet Loss Concealment output, since libopus conceals
    /// `frame_size` samples when given no input — passing the full max buffer
    /// would synthesize a 120ms burst instead of a single 20ms frame.
    frame_samples: usize,
}

impl OpusDecoder {
    /// Creates a new Opus decoder.
    /// Standard voice parameters: sample_rate = 48000, channels = Channels::Mono.
    /// `frame_samples` is the per-channel length of one transmitted frame (e.g. 960).
    pub fn new(
        sample_rate: u32,
        channels: Channels,
        frame_samples: usize,
    ) -> Result<Self, CodecError> {
        let decoder = Decoder::new(sample_rate, channels)
            .map_err(|e| CodecError::Init(format!("Failed to create Opus decoder: {}", e)))?;
        Ok(Self {
            decoder,
            channels,
            frame_samples,
        })
    }

    fn channels_count(&self) -> usize {
        match self.channels {
            Channels::Mono => 1,
            Channels::Stereo => 2,
        }
    }
}

impl AudioDecoder for OpusDecoder {
    fn decode(&mut self, compressed: Option<&[u8]>) -> Result<Vec<i16>, CodecError> {
        let channels_count = self.channels_count();

        let (mut pcm, input): (Vec<i16>, &[u8]) = match compressed {
            Some(data) if !data.is_empty() => {
                // Normal decode. Size the buffer to the maximum Opus frame (120ms =
                // 5760 samples/channel); libopus decodes the packet's true duration.
                (vec![0i16; 5760 * channels_count], data)
            }
            _ => {
                // Packet Loss Concealment: an empty input makes libopus synthesize
                // exactly `frame_samples` of concealment, so size the buffer to match.
                (vec![0i16; self.frame_samples * channels_count], &[])
            }
        };

        let decoded_per_channel = self
            .decoder
            .decode(input, &mut pcm, false)
            .map_err(|e| CodecError::Decode(format!("Opus decoding failed: {}", e)))?;

        pcm.truncate(decoded_per_channel * channels_count);
        Ok(pcm)
    }

    fn decode_fec(&mut self, next_payload: &[u8]) -> Result<Vec<i16>, CodecError> {
        let channels_count = self.channels_count();
        let mut pcm = vec![0i16; self.frame_samples * channels_count];

        let decoded_per_channel = self
            .decoder
            .decode(next_payload, &mut pcm, true)
            .map_err(|e| CodecError::Decode(format!("Opus FEC decoding failed: {}", e)))?;

        pcm.truncate(decoded_per_channel * channels_count);
        Ok(pcm)
    }
}

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

    #[test]
    fn test_opus_params_mapping() {
        let low = opus_params(AudioProfile::LowLatency);
        let bal = opus_params(AudioProfile::Balanced);
        let bad = opus_params(AudioProfile::BadNetwork);

        // LowLatency has no loss redundancy; the other two do.
        assert!(!low.inband_fec);
        assert_eq!(low.packet_loss_perc, 0);
        assert!(bal.inband_fec);
        assert!(bad.inband_fec);

        // Capture-side gating suppresses silence; no profile adds Opus DTX.
        assert!(!low.dtx && !bal.dtx && !bad.dtx);
        assert!(bad.packet_loss_perc > bal.packet_loss_perc);

        // BadNetwork trims base bitrate to make room for FEC redundancy.
        assert!(bad.bitrate < bal.bitrate);

        // All bitrates are sane positive voice rates.
        for p in [low, bal, bad] {
            assert!(p.bitrate > 0 && p.bitrate <= 64_000);
            assert!((0..=100).contains(&p.packet_loss_perc));
        }
    }

    #[test]
    fn test_apply_profile_sets_bitrate() {
        let mut encoder = OpusEncoder::new(48000, Channels::Mono, Application::Voip).unwrap();
        // Every profile applies cleanly to a real encoder...
        for profile in AudioProfile::ALL {
            encoder.apply_profile(profile).unwrap();
        }
        // ...and the last-applied bitrate is reflected by the encoder.
        encoder.apply_profile(AudioProfile::Balanced).unwrap();
        let want = opus_params(AudioProfile::Balanced).bitrate;
        assert_eq!(encoder.encoder.get_bitrate().unwrap(), Bitrate::Bits(want));
    }

    #[test]
    fn test_round_trip() {
        let mut encoder = OpusEncoder::new(48000, Channels::Mono, Application::Voip).unwrap();
        let mut decoder = OpusDecoder::new(48000, Channels::Mono, 960).unwrap();

        // 1. Round-trip shape: generate a 960-sample mono PCM frame (440Hz sine wave)
        let mut pcm = vec![0i16; 960];
        for (i, sample) in pcm.iter_mut().enumerate() {
            let t = (i as f32) / 48000.0;
            let val = (t * 440.0 * 2.0 * std::f32::consts::PI).sin();
            *sample = (val * 10000.0) as i16;
        }

        // encode it
        let compressed = encoder.encode(&pcm).unwrap();
        assert!(
            !compressed.is_empty(),
            "Compressed buffer should not be empty"
        );
        assert!(
            compressed.len() < pcm.len() * std::mem::size_of::<i16>(),
            "Compressed size ({}) should be smaller than raw PCM size ({})",
            compressed.len(),
            pcm.len() * std::mem::size_of::<i16>()
        );

        // decode it
        let decoded = decoder.decode(Some(&compressed)).unwrap();
        assert_eq!(
            decoded.len(),
            960,
            "Decoded sample count should be exactly 960"
        );

        // 2. Round-trip carries signal energy (not silence)
        let sum_sq: f64 = decoded.iter().map(|&x| (x as f64).powi(2)).sum();
        let rms = (sum_sq / decoded.len() as f64).sqrt();
        // Since input had amplitude ~10000, let's verify RMS is significantly above 0 (e.g. > 100.0)
        assert!(
            rms > 100.0,
            "Decoded signal should carry energy (RMS was {})",
            rms
        );
    }

    #[test]
    fn test_plc_sizing() {
        let mut decoder = OpusDecoder::new(48000, Channels::Mono, 960).unwrap();

        // decode(None) returns exactly frame_samples (960) samples
        let plc_none = decoder.decode(None).unwrap();
        assert_eq!(
            plc_none.len(),
            960,
            "decode(None) should yield exactly 960 samples"
        );

        // decode(Some(&[])) (empty slice) does the same
        let plc_empty = decoder.decode(Some(&[])).unwrap();
        assert_eq!(
            plc_empty.len(),
            960,
            "decode(Some(&[])) should yield exactly 960 samples"
        );
    }

    #[test]
    fn test_decoded_duration_follows_packet() {
        let mut encoder = OpusEncoder::new(48000, Channels::Mono, Application::Voip).unwrap();
        let mut decoder = OpusDecoder::new(48000, Channels::Mono, 960).unwrap();

        let pcm = vec![0i16; 960];
        let compressed = encoder.encode(&pcm).unwrap();
        let decoded = decoder.decode(Some(&compressed)).unwrap();
        assert_eq!(
            decoded.len(),
            960,
            "Decoded sample count should match packet duration"
        );
    }

    #[test]
    fn test_stereo_plc_sizing() {
        let mut decoder = OpusDecoder::new(48000, Channels::Stereo, 960).unwrap();

        // decode(None) returns exactly frame_samples * 2 (1920) samples
        let plc_none = decoder.decode(None).unwrap();
        assert_eq!(
            plc_none.len(),
            960 * 2,
            "Stereo decode(None) should yield exactly 1920 samples"
        );

        // decode(Some(&[])) (empty slice) does the same
        let plc_empty = decoder.decode(Some(&[])).unwrap();
        assert_eq!(
            plc_empty.len(),
            960 * 2,
            "Stereo decode(Some(&[])) should yield exactly 1920 samples"
        );
    }
}