bbr 0.1.0

use std::{
    collections::{HashMap, VecDeque},
    thread,
    time::{Duration, Instant},
};

use bbr::v3::{Bbr, BbrConfig, CwndEvent, RateSample, Transport};
use core_affinity;
use kanal::{Receiver, Sender, unbounded};
use rand::{prelude::*, rng};

/// Minimal packet representation for in-memory transport
#[derive(Debug, Clone)]
struct Packet {
    seq: u64,
    send_time: u64, // microseconds since start
}

/// ACK describing highest seq received and optional missing list
#[derive(Debug, Clone)]
struct AckPacket {
    ack_seq: u64,
    _send_time: u64,
    _recv_time: u64,
    missing: Vec<u64>,
}

/////////////////////////////////////////////////////////////////////////////////
/// In-memory sender side transport implementing the `Transport` trait
/////////////////////////////////////////////////////////////////////////////////
struct MemTransport {
    start: Instant,
    tx_data: Sender<Packet>,
    rx_ack: Receiver<AckPacket>,

    seq: u64,

    // Congestion-control stats
    delivered: u64,
    in_flight: u32,
    cwnd: u32,
    mss: u32,

    sent: HashMap<
        u64,
        (
            u64, /* send_time */
            u64, /* delivered_at_send */
            u32, /* in_flight */
        ),
    >,
    srtt_us: Option<u32>,
}

impl MemTransport {
    fn new(tx_data: Sender<Packet>, rx_ack: Receiver<AckPacket>) -> Self {
        Self {
            start: Instant::now(),
            tx_data,
            rx_ack,
            seq: 0,
            delivered: 0,
            in_flight: 0,
            cwnd: 10,
            mss: 1, // one packet == one MSS for the model
            sent: HashMap::new(),
            srtt_us: None,
        }
    }

    /// Try to transmit a packet according to cwnd
    fn maybe_send(&mut self) -> bool {
        if self.in_flight >= self.cwnd {
            return false;
        }
        let now = self.now_us();
        let pkt = Packet {
            seq: self.seq,
            send_time: now,
        };
        self.seq += 1;
        let _ = self.tx_data.send(pkt);
        self.sent
            .insert(self.seq - 1, (now, self.delivered, self.in_flight));
        self.in_flight += 1;
        true
    }

    /// Drain all available ACKs, returning whether any were processed
    fn process_acks(&mut self, bbr: &mut Bbr) -> bool {
        let mut any = false;
        loop {
            match self.rx_ack.try_recv() {
                Ok(Some(ack)) => {
                    any = true;
                    self.handle_ack(bbr, ack);
                }
                Ok(None) => break,
                Err(_) => break,
            }
        }
        any
    }

    fn handle_ack(&mut self, bbr: &mut Bbr, ack: AckPacket) {
        let now = self.now_us();
        // Newly ACKed list
        let mut newly_acked = Vec::new();
        let mut to_remove = Vec::new();
        for (&seq, &(send_time, deliv_at_send, inflight_at_send)) in &self.sent {
            if seq <= ack.ack_seq && !ack.missing.contains(&seq) {
                newly_acked.push((seq, send_time, deliv_at_send, inflight_at_send));
                to_remove.push(seq);
            }
        }
        for s in to_remove {
            self.sent.remove(&s);
            self.in_flight = self.in_flight.saturating_sub(1);
        }
        if newly_acked.is_empty() {
            return;
        }
        // RTT sample from earliest ACKed packet
        let rtt_us = (now - newly_acked[0].1) as i64;
        if rtt_us > 0 {
            self.srtt_us = Some(
                self.srtt_us
                    .map_or(rtt_us as u32, |srtt| (7 * srtt + rtt_us as u32) / 8),
            );
        }

        // Update global delivered counter *before* building the rate sample so that it reflects
        // bytes delivered up to this ACK.
        self.delivered += newly_acked.len() as u64;

        // The BBR spec defines delivered as the total packets delivered since the first packet
        // in the sample was transmitted.
        let first_sent = &newly_acked[0];
        let delivered_pkts = self.delivered.saturating_sub(first_sent.2); // total delivered since that send
        let interval_us = (now - first_sent.1).max(1); // elapsed time since that packet was sent

        let rs = RateSample {
            delivered: delivered_pkts,
            interval_us,
            rtt_us,
            losses: 0,
            acked_sacked: newly_acked.len() as u32,
            prior_in_flight: first_sent.3,
            is_ack_delayed: false,
            is_app_limited: false,
            tx_in_flight: first_sent.3,
            lost: 0,
            delivered_ce: 0,
            prior_delivered: first_sent.2,
            is_acking_tlp_retrans_seq: false,
        };
        bbr.update(self, &rs);
    }
}

impl Transport for MemTransport {
    fn now_us(&self) -> u64 {
        self.start.elapsed().as_micros() as u64
    }
    fn delivered(&self) -> u64 {
        self.delivered
    }
    fn delivered_ce(&self) -> u64 {
        0
    }
    fn packets_in_flight(&self) -> u32 {
        self.in_flight
    }
    fn cwnd(&self) -> u32 {
        self.cwnd
    }
    fn set_cwnd(&mut self, cwnd: u32) {
        self.cwnd = cwnd.max(self.in_flight + 2).max(4);
    }
    fn mss(&self) -> u32 {
        self.mss
    }
    fn pacing_rate(&self) -> u64 {
        0
    }
    fn set_pacing_rate(&mut self, _rate: u64) {}
    fn max_pacing_rate(&self) -> u64 {
        u64::MAX
    }
    fn ecn_eligible(&self) -> bool {
        false
    }
    fn is_cwnd_limited(&self) -> bool {
        self.in_flight >= self.cwnd.saturating_sub(3)
    }
    fn random_below(&self, max: u32) -> u32 {
        if max == 0 {
            0
        } else {
            rng().random_range(0..max)
        }
    }
    fn srtt_us(&self) -> Option<u32> {
        self.srtt_us
    }
    fn rcv_nxt(&self) -> u32 {
        self.delivered as u32
    }
    fn in_recovery(&self) -> bool {
        false
    }
    fn lost_out(&self) -> u32 {
        0
    }
    fn set_fast_ack_mode(&mut self, _enabled: bool) {}
}

/////////////////////////////////////////////////////////////////////////////////
/// Receiver thread – optional loss simulation
/////////////////////////////////////////////////////////////////////////////////
fn receiver_task(
    loss_pct: f64,
    rx: Receiver<Packet>,
    tx_ack: Sender<AckPacket>,
    core_idx: Option<usize>,
) {
    if let Some(idx) = core_idx {
        if let Some(core) = core_affinity::get_core_ids().and_then(|v| v.get(idx).cloned()) {
            core_affinity::set_for_current(core);
        }
    }
    let mut highest_seq: i64 = -1;
    let mut missing: Vec<u64> = Vec::new();
    let mut rng = rng();

    const ACK_EVERY: u32 = 32;
    let mut batch: u32 = 0;
    let mut last_ack_time = Instant::now();

    while let Ok(pkt) = rx.recv() {
        // Simulate drop
        if rng.random_bool(loss_pct / 100.0) {
            continue;
        }
        let now = pkt.send_time + 1000; // fixed 1-ms propagation for rtt realism

        if highest_seq < 0 || pkt.seq as i64 > highest_seq {
            for seq in (highest_seq + 1) as u64..pkt.seq {
                missing.push(seq);
                if missing.len() > 128 {
                    missing.remove(0);
                }
            }
            highest_seq = pkt.seq as i64;
        } else if let Some(pos) = missing.iter().position(|&s| s == pkt.seq) {
            missing.remove(pos);
        }

        batch += 1;
        if batch >= ACK_EVERY || last_ack_time.elapsed() >= Duration::from_millis(2) {
            let ack = AckPacket {
                ack_seq: highest_seq as u64,
                _send_time: pkt.send_time,
                _recv_time: now,
                missing: missing.clone(),
            };
            tx_ack.send(ack).ok();
            batch = 0;
            last_ack_time = Instant::now();
        }
    }
}

/////////////////////////////////////////////////////////////////////////////////
/// Router task: adds fixed latency and finite queue (drop-tail)
/////////////////////////////////////////////////////////////////////////////////
fn router_task(
    latency_us: u64,
    rate_mbps: f64,
    rx_in: Receiver<Packet>,
    tx_out: Sender<Packet>,
    core_idx: Option<usize>,
) {
    if let Some(idx) = core_idx {
        if let Some(core) = core_affinity::get_core_ids().and_then(|v| v.get(idx).cloned()) {
            core_affinity::set_for_current(core);
        }
    }
    let start = Instant::now();
    let mut pending: VecDeque<(Packet, u64)> = VecDeque::new(); // (packet, due_time_ns)

    // interval_ns = (pkt_bits / rate_bps) * 1e9
    let pkt_bits: f64 = 1500.0 * 8.0;
    let interval_ns: u64 = if rate_mbps <= 0.0 {
        0
    } else {
        ((pkt_bits * 1000.0) / rate_mbps) as u64 // bits*1e9 / (rate*1e6) == bits*1000 / rate
    };

    let mut next_send_time: u64 = 0; // in nanoseconds since start

    let now_ns = |s: &Instant| s.elapsed().as_nanos() as u64;

    // Stats for packets-per-second reporting
    let mut delivered_pkts: u64 = 0;
    let mut last_report = Instant::now();

    loop {
        let now = now_ns(&start);

        // deliver packets whose delivery time has arrived
        while let Some((pkt, due)) = pending.front().cloned() {
            if due <= now {
                let _ = tx_out.send(pkt);
                pending.pop_front();

                // stats
                delivered_pkts += 1;
            } else {
                break;
            }
        }

        match rx_in.try_recv() {
            Ok(Some(pkt)) => {
                // schedule send respecting bandwidth limiter (nanosecond precision)
                let send_time = if next_send_time < now {
                    now
                } else {
                    next_send_time
                };
                next_send_time = send_time.saturating_add(interval_ns.max(1));

                // Convert latency (provided in microseconds) to nanoseconds when
                // computing the delivery time.
                let deliver_time = send_time + latency_us * 1000;
                pending.push_back((pkt, deliver_time));
            }
            Ok(None) => {
                thread::sleep(Duration::from_micros(50));
            }
            Err(_) => break,
        }

        // Periodic stats output
        if last_report.elapsed() >= Duration::from_secs(1) {
            println!(
                "router: {:.2} kpps (queue={})",
                delivered_pkts as f64 / 1000.0 / last_report.elapsed().as_secs_f64(),
                pending.len()
            );
            delivered_pkts = 0;
            last_report = Instant::now();
        }
    }
}

/////////////////////////////////////////////////////////////////////////////////
/// Main – spawns sender & receiver and prints stats
/////////////////////////////////////////////////////////////////////////////////
fn main() {
    let args: Vec<String> = std::env::args().collect();
    if args.len() != 4 {
        eprintln!(
            "Usage: {} <loss-%> <latency-ms> <rate-Mbps>\n  Example: {} 0 10 1000",
            args.get(0).map(|s| s.as_str()).unwrap_or("mem_transport"),
            args.get(0).map(|s| s.as_str()).unwrap_or("mem_transport")
        );
        std::process::exit(1);
    }
    let loss_pct: f64 = args[1].parse().expect("loss % must be a number");
    let latency_ms: u64 = args[2].parse().expect("latency-ms must be an integer");
    let rate_mbps: f64 = args[3].parse().expect("rate-Mbps must be a number");

    let (tx_data, rx_data_sender) = unbounded();
    let (tx_router_out, rx_data_receiver) = unbounded();
    let (tx_ack, rx_ack) = unbounded(); // receiver to sender

    // Spawn router
    thread::spawn(move || {
        router_task(
            latency_ms * 1000,
            rate_mbps,
            rx_data_sender,
            tx_router_out,
            Some(1),
        )
    });
    // Spawn receiver
    thread::spawn(move || receiver_task(loss_pct, rx_data_receiver, tx_ack, Some(2)));

    if let Some(core) = core_affinity::get_core_ids().and_then(|v| v.get(0).cloned()) {
        core_affinity::set_for_current(core);
    }

    // Sender path
    let mut transport = MemTransport::new(tx_data, rx_ack);
    let mut bbr = Bbr::new(BbrConfig::default());
    bbr.init(&mut transport);
    bbr.cwnd_event(&mut transport, CwndEvent::TxStart);

    let mut last_report = Instant::now();
    let mut last_delivered = 0u64;
    const BYTES_PER_PACKET: u64 = 1500;

    loop {
        // Drive ACK processing
        let _ = transport.process_acks(&mut bbr);
        // try to send as much as cwnd allows
        while transport.maybe_send() {}

        if last_report.elapsed() >= Duration::from_secs(1) {
            let delivered_pkts = transport.delivered - last_delivered;
            last_delivered = transport.delivered;
            let throughput_mbps = (delivered_pkts * BYTES_PER_PACKET * 8) as f64
                / 1_000_000.0
                / last_report.elapsed().as_secs_f64();
            println!(
                "stats: cwnd={} inflight={} delivered={} rtt={:?} throughput={:.2} Mbps",
                transport.cwnd,
                transport.in_flight,
                transport.delivered,
                std::time::Duration::from_micros(transport.srtt_us.unwrap_or(0) as u64),
                throughput_mbps
            );
            last_report = Instant::now();
        }
        std::hint::spin_loop();
    }
}