btest-rs/src/bandwidth.rs

use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64};
use std::sync::Arc;
use std::time::Duration;

/// Shared state for bandwidth tracking between TX/RX threads and status reporter.
#[derive(Debug)]
pub struct BandwidthState {
    pub tx_bytes: AtomicU64,
    pub rx_bytes: AtomicU64,
    pub tx_speed: AtomicU32,
    pub tx_speed_changed: AtomicBool,
    pub running: AtomicBool,
    pub rx_packets: AtomicU64,
    pub rx_lost_packets: AtomicU64,
    pub last_udp_seq: AtomicU32,
}

impl BandwidthState {
    pub fn new() -> Arc<Self> {
        Arc::new(Self {
            tx_bytes: AtomicU64::new(0),
            rx_bytes: AtomicU64::new(0),
            tx_speed: AtomicU32::new(0),
            tx_speed_changed: AtomicBool::new(false),
            running: AtomicBool::new(true),
            rx_packets: AtomicU64::new(0),
            rx_lost_packets: AtomicU64::new(0),
            last_udp_seq: AtomicU32::new(0),
        })
    }
}

/// Calculate the sleep interval between packets to achieve target bandwidth.
/// Returns None if speed is unlimited (0).
pub fn calc_send_interval(tx_speed_bps: u32, tx_size: u16) -> Option<Duration> {
    if tx_speed_bps == 0 {
        return None;
    }

    let bits_per_packet = tx_size as u64 * 8;
    let interval_ns = (1_000_000_000u64 * bits_per_packet) / tx_speed_bps as u64;

    // Replicate MikroTik behavior: if interval > 500ms, clamp to 1 second
    if interval_ns > 500_000_000 {
        Some(Duration::from_secs(1))
    } else {
        Some(Duration::from_nanos(interval_ns.max(1)))
    }
}

/// Format a bandwidth value in human-readable form.
pub fn format_bandwidth(bits_per_sec: f64) -> String {
    if bits_per_sec >= 1_000_000_000.0 {
        format!("{:.2} Gbps", bits_per_sec / 1_000_000_000.0)
    } else if bits_per_sec >= 1_000_000.0 {
        format!("{:.2} Mbps", bits_per_sec / 1_000_000.0)
    } else if bits_per_sec >= 1_000.0 {
        format!("{:.2} Kbps", bits_per_sec / 1_000.0)
    } else {
        format!("{:.0} bps", bits_per_sec)
    }
}

/// Parse bandwidth string like "100M", "1G", "500K", "1000000"
pub fn parse_bandwidth(s: &str) -> std::result::Result<u32, anyhow::Error> {
    let s = s.trim();
    if s.is_empty() {
        return Err(anyhow::anyhow!("Empty bandwidth string"));
    }

    let (num_str, multiplier) = match s.as_bytes().last() {
        Some(b'G' | b'g') => (&s[..s.len() - 1], 1_000_000_000u64),
        Some(b'M' | b'm') => (&s[..s.len() - 1], 1_000_000u64),
        Some(b'K' | b'k') => (&s[..s.len() - 1], 1_000u64),
        _ => (s, 1u64),
    };

    let num: f64 = num_str
        .parse()
        .map_err(|e| anyhow::anyhow!("Invalid bandwidth number '{}': {}", num_str, e))?;
    let result = (num * multiplier as f64) as u64;
    if result > u32::MAX as u64 {
        Err(anyhow::anyhow!("Bandwidth {} exceeds maximum (4 Gbps)", s))
    } else {
        Ok(result as u32)
    }
}

/// Print a status line for a reporting interval.
pub fn print_status(
    interval_num: u32,
    direction: &str,
    bytes: u64,
    elapsed: Duration,
    lost_packets: Option<u64>,
) {
    let secs = elapsed.as_secs_f64();
    let bits = bytes as f64 * 8.0;
    let bw = if secs > 0.0 { bits / secs } else { 0.0 };

    let loss_str = match lost_packets {
        Some(lost) if lost > 0 => format!("  lost: {}", lost),
        _ => String::new(),
    };

    println!(
        "[{:4}] {:>3}  {} ({} bytes){}",
        interval_num,
        direction,
        format_bandwidth(bw),
        bytes,
        loss_str,
    );
}

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

    #[test]
    fn test_parse_bandwidth() {
        assert_eq!(parse_bandwidth("100M").unwrap(), 100_000_000);
        assert_eq!(parse_bandwidth("1G").unwrap(), 1_000_000_000);
        assert_eq!(parse_bandwidth("500K").unwrap(), 500_000);
        assert_eq!(parse_bandwidth("1000000").unwrap(), 1_000_000);
        assert_eq!(parse_bandwidth("1.5M").unwrap(), 1_500_000);
    }

    #[test]
    fn test_calc_interval() {
        // 100Mbps with 1500 byte packets
        let interval = calc_send_interval(100_000_000, 1500).unwrap();
        // Expected: (1e9 * 1500 * 8) / 100_000_000 = 120_000 ns = 120 us
        assert_eq!(interval.as_nanos(), 120_000);

        // Unlimited
        assert!(calc_send_interval(0, 1500).is_none());
    }

    #[test]
    fn test_format_bandwidth() {
        assert_eq!(format_bandwidth(100_000_000.0), "100.00 Mbps");
        assert_eq!(format_bandwidth(1_500_000_000.0), "1.50 Gbps");
        assert_eq!(format_bandwidth(500_000.0), "500.00 Kbps");
    }
}