use std::{
    env,
    fs::File,
    io::{self, Read, Seek, SeekFrom},
    thread::sleep,
    time::{Duration, Instant},
};

struct LinkSpeed {
    rx: File,
    tx: File,
    rx_bytes: u64,
    tx_bytes: u64,
    time: Instant,
}

impl LinkSpeed {
    pub fn new(dev: String) -> io::Result<Self> {
        let rx_bytes_file = File::open(format!("/sys/class/net/{dev}/statistics/rx_bytes"))?;
        let tx_bytes_file = File::open(format!("/sys/class/net/{dev}/statistics/tx_bytes"))?;

        Ok(Self {
            rx: rx_bytes_file,
            tx: tx_bytes_file,
            rx_bytes: 0,
            tx_bytes: 0,
            time: Instant::now(),
        })
    }

    fn read_bytes(file: &mut File) -> io::Result<u64> {
        file.seek(SeekFrom::Start(0))?;
        let mut buf = String::new();
        file.read_to_string(&mut buf)?;
        Ok(buf.trim().parse().unwrap_or(0))
    }

    fn update(&mut self) -> io::Result<(u64, u64)> {
        let rx_bytes = Self::read_bytes(&mut self.rx)?;
        let tx_bytes = Self::read_bytes(&mut self.tx)?;

        let mut rx_diff = rx_bytes - self.rx_bytes;
        let mut tx_diff = tx_bytes - self.tx_bytes;

        if self.rx_bytes == 0 {
            self.rx_bytes = rx_bytes;
            rx_diff = 0;
        }

        if self.tx_bytes == 0 {
            self.tx_bytes = tx_bytes;
            tx_diff = 0;
        }

        self.rx_bytes = rx_bytes;
        self.tx_bytes = tx_bytes;

        Ok((rx_diff, tx_diff))
    }

    pub fn get_measurement(&mut self) -> (f64, f64) {
        let elapsed = self.time.elapsed().as_secs_f64();
        let (rx, tx) = self.update().unwrap();
        let rx_speed = rx as f64 / elapsed;
        let tx_speed = tx as f64 / elapsed;
        self.time = Instant::now();
        (rx_speed, tx_speed)
    }
}

impl Iterator for LinkSpeed {
    type Item = (f64, f64);

    fn next(&mut self) -> Option<Self::Item> {
        Some(self.get_measurement())
    }
}

fn main() {
    let netdev_name = env::args().nth(1).expect("No network device provided");

    let link_speed = LinkSpeed::new(netdev_name).expect("Failed to create LinkSpeed object");

    link_speed.for_each(|(rx_speed, tx_speed)| {
        println!(
            "RX: {:.0} MBit/s, TX: {:.0} MBit/s",
            rx_speed / 1024.0 / 1024.0 * 8.0,
            tx_speed / 1024.0 / 1024.0 * 8.0
        );
        sleep(Duration::from_millis(1000));
    });
}