msf-rtp 0.4.1

use std::time::Duration;

use bytes::{Buf, BufMut, Bytes, BytesMut};
use zerocopy::{
    byteorder::network_endian::{U32, U64},
    FromBytes, Immutable, IntoBytes, KnownLayout, SizeError, Unaligned,
};

use crate::{InvalidInput, RtcpPacket, RtcpPacketType};

/// Helper struct.
#[derive(Copy, Clone, KnownLayout, Immutable, Unaligned, IntoBytes, FromBytes)]
#[repr(C)]
struct RawReportBlock {
    ssrc: U32,
    loss: U32,
    extended_sequence_number: U32,
    jitter: U32,
    last_sr_timestamp: U32,
    delay_since_last_sr: U32,
}

/// Sender/receiver report block.
#[derive(Copy, Clone)]
pub struct ReportBlock {
    ssrc: u32,
    loss: u32,
    extended_sequence_number: u32,
    jitter: u32,
    last_sr_timestamp: u32,
    delay_since_last_sr: u32,
}

impl ReportBlock {
    /// Size of an encoded report block.
    pub const RAW_SIZE: usize = std::mem::size_of::<RawReportBlock>();

    /// Create a new report block.
    #[inline]
    pub const fn new(ssrc: u32) -> Self {
        Self {
            ssrc,
            loss: 0,
            extended_sequence_number: 0,
            jitter: 0,
            last_sr_timestamp: 0,
            delay_since_last_sr: 0,
        }
    }

    /// Decode a report block from given data.
    pub fn decode(data: &mut Bytes) -> Result<Self, InvalidInput> {
        let (raw, _) = RawReportBlock::ref_from_prefix(data)
            .map_err(SizeError::from)
            .map_err(|_| InvalidInput::new())?;

        let res = Self {
            ssrc: raw.ssrc.get(),
            loss: raw.loss.get(),
            extended_sequence_number: raw.extended_sequence_number.get(),
            jitter: raw.jitter.get(),
            last_sr_timestamp: raw.last_sr_timestamp.get(),
            delay_since_last_sr: raw.delay_since_last_sr.get(),
        };

        data.advance(std::mem::size_of::<RawReportBlock>());

        Ok(res)
    }

    /// Encode the report block.
    pub fn encode(&self, buf: &mut BytesMut) {
        let raw = RawReportBlock {
            ssrc: U32::new(self.ssrc),
            loss: U32::new(self.loss),
            extended_sequence_number: U32::new(self.extended_sequence_number),
            jitter: U32::new(self.jitter),
            last_sr_timestamp: U32::new(self.last_sr_timestamp),
            delay_since_last_sr: U32::new(self.delay_since_last_sr),
        };

        buf.extend_from_slice(raw.as_bytes());
    }

    /// Get SSRC.
    #[inline]
    pub fn ssrc(&self) -> u32 {
        self.ssrc
    }

    /// Set SSRC.
    #[inline]
    pub fn with_ssrc(mut self, ssrc: u32) -> Self {
        self.ssrc = ssrc;
        self
    }

    /// Get fractional loss as 0.8 fixed point number.
    #[inline]
    pub fn fractional_loss(&self) -> u8 {
        (self.loss >> 24) as u8
    }

    /// Set fractional loss as 0.8 fixed point number.
    #[inline]
    pub fn with_fractional_loss(mut self, loss: u8) -> Self {
        self.loss &= 0x00ffffff;
        self.loss |= (loss as u32) << 24;
        self
    }

    /// Get cumulative packet loss (the precision is only up to 24 bits).
    #[inline]
    pub fn cumulative_loss(&self) -> i32 {
        ((self.loss << 8) as i32) >> 8
    }

    /// Set cumulative packet loss (the precision is only up to 24 bits).
    #[inline]
    pub fn with_cumulative_loss(mut self, loss: i32) -> Self {
        let min = -(1i32 << 23);
        let max = (1i32 << 23) - 1;

        let loss = loss.clamp(min, max) as u32;

        self.loss &= 0xff000000;
        self.loss |= loss & 0x00ffffff;
        self
    }

    /// Set loss calculated from a given number of expected packets vs. a given
    /// number of received packets.
    pub fn with_loss(self, expected: u64, received: u64) -> Self {
        let delta = expected as i64 - received as i64;

        let fraction = if delta < 0 {
            0
        } else {
            ((delta as u64) << 8) / expected
        };

        self.with_fractional_loss(fraction as u8)
            .with_cumulative_loss(delta as i32)
    }

    /// Get extended highest sequence number.
    #[inline]
    pub fn extended_sequence_number(&self) -> u32 {
        self.extended_sequence_number
    }

    /// Set the extended sequence number.
    #[inline]
    pub fn with_extended_sequence_number(mut self, n: u32) -> Self {
        self.extended_sequence_number = n;
        self
    }

    /// Get jitter.
    #[inline]
    pub fn jitter(&self) -> u32 {
        self.jitter
    }

    /// Set the jitter.
    #[inline]
    pub fn with_jitter(mut self, jitter: u32) -> Self {
        self.jitter = jitter;
        self
    }

    /// Get NTP timestamp of the last sender report (after truncating to the
    /// middle 32 bits).
    ///
    /// The returned timestamp is a 32.32 fixed point number.
    #[inline]
    pub fn last_sr_timestamp(&self) -> u64 {
        (self.last_sr_timestamp as u64) << 16
    }

    /// Set NTP timestamp of the last sender report.
    ///
    /// The timestamp is expected to be a 32.32 fixed point number and it will
    /// be truncated to the middle 32 bits.
    #[inline]
    pub fn with_last_sr_timestamp(mut self, ts: u64) -> Self {
        self.last_sr_timestamp = (ts >> 16) as u32;
        self
    }

    /// Get delay since the last sender report.
    #[inline]
    pub fn delay_since_last_sr(&self) -> Duration {
        let secs = (self.delay_since_last_sr >> 16) as u64;
        let nanos = ((self.delay_since_last_sr & 0xffff) as u64 * 1_000_000_000) >> 16;

        Duration::new(secs, nanos as u32)
    }

    /// Set delay since the last sender report.
    #[inline]
    pub fn with_delay_since_last_sr(mut self, delay: Duration) -> Self {
        let secs = (delay.as_secs() << 16) as u32;
        let fraction = (((delay.subsec_nanos() as u64) << 16) / 1_000_000_000) as u32;

        self.delay_since_last_sr = secs + fraction;
        self
    }

    /// Get size of the encoded report block.
    #[inline]
    pub fn raw_size(&self) -> usize {
        Self::RAW_SIZE
    }
}

/// Helper struct.
#[derive(Copy, Clone, KnownLayout, Immutable, Unaligned, IntoBytes, FromBytes)]
#[repr(C)]
struct RawSenderReportHeader {
    sender_ssrc: U32,
    ntp_timestamp: U64,
    rtp_timestamp: U32,
    packet_count: U32,
    octet_count: U32,
}

/// Sender report.
#[derive(Clone)]
pub struct SenderReport {
    sender_ssrc: u32,
    ntp_timestamp: u64,
    rtp_timestamp: u32,
    packet_count: u32,
    octet_count: u32,
    report_blocks: Vec<ReportBlock>,
}

impl SenderReport {
    /// Create a new sender report.
    #[inline]
    pub const fn new(sender_ssrc: u32) -> Self {
        Self {
            sender_ssrc,
            ntp_timestamp: 0,
            rtp_timestamp: 0,
            packet_count: 0,
            octet_count: 0,
            report_blocks: Vec::new(),
        }
    }

    /// Decode sender report.
    pub fn decode(packet: &RtcpPacket) -> Result<Self, InvalidInput> {
        let header = packet.header();

        let mut data = packet.stripped_payload();

        let (raw, _) = RawSenderReportHeader::ref_from_prefix(&data)
            .map_err(SizeError::from)
            .map_err(|_| InvalidInput::new())?;

        let mut res = Self {
            sender_ssrc: raw.sender_ssrc.get(),
            ntp_timestamp: raw.ntp_timestamp.get(),
            rtp_timestamp: raw.rtp_timestamp.get(),
            packet_count: raw.packet_count.get(),
            octet_count: raw.octet_count.get(),
            report_blocks: Vec::with_capacity(header.item_count() as usize),
        };

        data.advance(std::mem::size_of::<RawSenderReportHeader>());

        for _ in 0..header.item_count() {
            res.report_blocks.push(ReportBlock::decode(&mut data)?);
        }

        Ok(res)
    }

    /// Encode the sender report.
    pub fn encode(&self) -> RtcpPacket {
        let mut payload = BytesMut::with_capacity(self.raw_size());

        let raw = RawSenderReportHeader {
            sender_ssrc: U32::new(self.sender_ssrc),
            ntp_timestamp: U64::new(self.ntp_timestamp),
            rtp_timestamp: U32::new(self.rtp_timestamp),
            packet_count: U32::new(self.packet_count),
            octet_count: U32::new(self.octet_count),
        };

        payload.extend_from_slice(raw.as_bytes());

        for block in &self.report_blocks {
            block.encode(&mut payload);
        }

        RtcpPacket::new(RtcpPacketType::SR)
            .with_item_count(self.report_blocks.len() as u8)
            .with_payload(payload.freeze(), 0)
    }

    /// Get SSRC identifier of the sender.
    #[inline]
    pub fn sender_ssrc(&self) -> u32 {
        self.sender_ssrc
    }

    /// Set the SSRC identifier of the sender.
    #[inline]
    pub fn with_sender_ssrc(mut self, ssrc: u32) -> Self {
        self.sender_ssrc = ssrc;
        self
    }

    /// Get NTP timestamp as a 32.32 fixed point number.
    #[inline]
    pub fn ntp_timestamp(&self) -> u64 {
        self.ntp_timestamp
    }

    /// Set the NTP timestamp as a 32.32. fixed point number.
    #[inline]
    pub fn with_ntp_timestamp(mut self, timestamp: u64) -> Self {
        self.ntp_timestamp = timestamp;
        self
    }

    /// Get RTP timestamp.
    #[inline]
    pub fn rtp_timestamp(&self) -> u32 {
        self.rtp_timestamp
    }

    /// Set the RTP timestamp.
    #[inline]
    pub fn with_rtp_timestamp(mut self, timestamp: u32) -> Self {
        self.rtp_timestamp = timestamp;
        self
    }

    /// Get packet count.
    #[inline]
    pub fn packet_count(&self) -> u32 {
        self.packet_count
    }

    /// Set the packet count.
    #[inline]
    pub fn with_packet_count(mut self, count: u32) -> Self {
        self.packet_count = count;
        self
    }

    /// Get octet count.
    #[inline]
    pub fn octet_count(&self) -> u32 {
        self.octet_count
    }

    /// Set the octet count.
    #[inline]
    pub fn with_octet_count(mut self, count: u32) -> Self {
        self.octet_count = count;
        self
    }

    /// Get report blocks.
    #[inline]
    pub fn report_blocks(&self) -> &[ReportBlock] {
        &self.report_blocks
    }

    /// Set the report blocks.
    ///
    /// # Panics
    /// The method will panic if the number of report blocks is greater than
    /// 31.
    #[inline]
    pub fn with_report_blocks<T>(mut self, blocks: T) -> Self
    where
        T: Into<Vec<ReportBlock>>,
    {
        let blocks = blocks.into();

        assert!(blocks.len() < 32);

        self.report_blocks = blocks;
        self
    }

    /// Get size of the encoded sender report.
    #[inline]
    pub fn raw_size(&self) -> usize {
        std::mem::size_of::<RawSenderReportHeader>()
            + std::mem::size_of::<RawReportBlock>() * self.report_blocks.len()
    }
}

/// Receiver report.
#[derive(Clone)]
pub struct ReceiverReport {
    sender_ssrc: u32,
    report_blocks: Vec<ReportBlock>,
}

impl ReceiverReport {
    /// Create a new receiver report.
    #[inline]
    pub const fn new(sender_ssrc: u32) -> Self {
        Self {
            sender_ssrc,
            report_blocks: Vec::new(),
        }
    }

    /// Decode receiver report.
    pub fn decode(packet: &RtcpPacket) -> Result<Self, InvalidInput> {
        let header = packet.header();

        let mut data = packet.stripped_payload();

        if data.len() < 4 {
            return Err(InvalidInput::new());
        }

        let mut res = Self {
            sender_ssrc: data.get_u32(),
            report_blocks: Vec::with_capacity(header.item_count() as usize),
        };

        for _ in 0..header.item_count() {
            res.report_blocks.push(ReportBlock::decode(&mut data)?);
        }

        Ok(res)
    }

    /// Encode the sender report.
    pub fn encode(&self) -> RtcpPacket {
        let mut payload = BytesMut::with_capacity(self.raw_size());

        payload.put_u32(self.sender_ssrc);

        for block in &self.report_blocks {
            block.encode(&mut payload);
        }

        RtcpPacket::new(RtcpPacketType::RR)
            .with_item_count(self.report_blocks.len() as u8)
            .with_payload(payload.freeze(), 0)
    }

    /// Get SSRC identifier of the sender.
    #[inline]
    pub fn sender_ssrc(&self) -> u32 {
        self.sender_ssrc
    }

    /// Set the SSRC identifier of the sender.
    #[inline]
    pub fn with_sender_ssrc(mut self, ssrc: u32) -> Self {
        self.sender_ssrc = ssrc;
        self
    }

    /// Get report blocks.
    #[inline]
    pub fn report_blocks(&self) -> &[ReportBlock] {
        &self.report_blocks
    }

    /// Set the report blocks.
    ///
    /// # Panics
    /// The method will panic if the number of report blocks is greater than
    /// 31.
    #[inline]
    pub fn with_report_blocks<T>(mut self, blocks: T) -> Self
    where
        T: Into<Vec<ReportBlock>>,
    {
        let blocks = blocks.into();

        assert!(blocks.len() < 32);

        self.report_blocks = blocks;
        self
    }

    /// Get size of the encoded sender report.
    #[inline]
    pub fn raw_size(&self) -> usize {
        4 + std::mem::size_of::<RawReportBlock>() * self.report_blocks.len()
    }
}