use crate::{Pid, jiffy::Jiffies};

/// Per-process NIC bandwidth configuration.
#[derive(Clone, Copy, Default)]
pub enum BandwidthConfig {
    /// No bandwidth limit (only network links latency).
    #[default]
    Unbounded,

    /// Limits throughput to the given number of bytes per jiffy.
    Bounded { inbound: usize, outbound: usize },
}

pub(crate) struct Bandwidth {
    in_bandwidth: usize,
    out_bandwidth: usize,
    in_pased: Vec<usize>,
    out_pased: Vec<usize>,
}

impl Bandwidth {
    pub(crate) fn new(bandwidth_type: BandwidthConfig, proc_num: usize) -> Self {
        let (in_bandwidth, out_bandwidth) = match bandwidth_type {
            BandwidthConfig::Unbounded => (usize::MAX, usize::MAX),
            BandwidthConfig::Bounded { inbound, outbound } => (inbound, outbound),
        };

        Self {
            in_bandwidth,
            out_bandwidth,
            in_pased: vec![0; proc_num],
            out_pased: vec![0; proc_num],
        }
    }

    pub(crate) fn try_recv(
        &mut self,
        message_size: usize,
        target: Pid,
        now: Jiffies,
    ) -> Option<Jiffies> {
        if self.in_bandwidth == usize::MAX {
            return None;
        }

        self.in_pased[target] += message_size;

        if self.in_pased[target] > now.0 * self.in_bandwidth {
            return Some(Jiffies(self.in_pased[target] / self.in_bandwidth));
        }

        None
    }

    pub(crate) fn try_send(
        &mut self,
        message_size: usize,
        target: Pid,
        now: Jiffies,
    ) -> Option<Jiffies> {
        if self.out_bandwidth == usize::MAX {
            return None;
        }

        // We need to immediately add message size to the passed count. Why?
        // Imagine there is several network messasages sent at 5 Jiffies, without
        // immediate accounting of message size they will observe same bandwidth contraint
        // without observing each other that can lead to wrong extra latency computation.
        // The same holds fot recv bandwidth constraint.
        self.out_pased[target] += message_size;

        if self.out_pased[target] > now.0 * self.out_bandwidth {
            return Some(Jiffies(self.out_pased[target] / self.out_bandwidth));
        }

        None
    }
}
