i3xs/src/widgets/cpu.rs

use std::collections::vec_deque::VecDeque;
use std::fs::File;
use std::io::{BufRead, BufReader, Result};

use i3monkit::{Block, ColorRGB, Widget, WidgetUpdate};

use crate::spark;

#[derive(Debug)]
struct Stat {
    usage: u64,
    total: u64,
}

/// The CPU usage widget
///
/// This widget draws a CPU usage pertentage bar on your i3 status bar.
pub struct CpuWidget {
    last_stats: Vec<Stat>,
    history: Vec<VecDeque<f32>>,
}

fn color_lerp(c1: ColorRGB, c2: ColorRGB, alpha: f32) -> ColorRGB {
    ColorRGB(
        (c1.0 as f32 * (1. - alpha) + c2.0 as f32 * alpha) as u8,
        (c1.1 as f32 * (1. - alpha) + c2.1 as f32 * alpha) as u8,
        (c1.2 as f32 * (1. - alpha) + c2.2 as f32 * alpha) as u8,
    )
}

fn color_to_string(c: ColorRGB) -> String {
    format!("#{:02x}{:02x}{:02x}", c.0, c.1, c.2)
}

impl CpuWidget {
    fn read_status() -> Result<Vec<Stat>> {
        let file = File::open("/proc/stat")?;
        let reader = BufReader::new(file);
        let mut stats: Vec<_> = Vec::new();
        for line in reader.lines() {
            let line = line?;
            let tokens: Vec<_> = line.trim().split(|c| c == ' ' || c == '\t').collect();
            // The sum total CPU entry has a space instead of a number, so we skip the line when
            // there's an extra empty column.
            if tokens[1] == "" {
                continue;
            }
            if tokens[0].starts_with("cpu") {
                let parsed = tokens[1..5]
                    .iter()
                    .map(|x| u64::from_str_radix(x, 10).unwrap())
                    .collect::<Vec<_>>();
                // [0] user
                // [1] nice
                // [2] system
                // [3] idle
                let usage = parsed[0] + parsed[1] + parsed[2];
                let total = usage + parsed[3];
                stats.push(Stat { usage, total })
            }
        }
        Ok(stats)
    }

    pub fn new() -> Self {
        let last_stats = Self::read_status().unwrap();
        let num_samples = 1;
        let history = (0..last_stats.len())
            .map(|_| (0..num_samples).map(|_| 0.).collect::<VecDeque<f32>>())
            .collect::<Vec<VecDeque<f32>>>();
        let ret = Self {
            last_stats,
            history,
        };

        return ret;
    }

    fn draw_history(&mut self) -> Option<String> {
        let stats = Self::read_status().ok()?;
        let percentages: Vec<_> = stats
            .iter()
            .zip(self.last_stats.iter())
            .map(|(cur, old)| {
                if cur.total == old.total {
                    return 0.;
                }
                (cur.usage - old.usage) as f32 / (cur.total - old.total) as f32
            })
            .collect();
        for (i, p) in percentages.iter().enumerate() {
            self.history[i].push_back(*p);
            self.history[i].pop_front();
        }
        self.last_stats = stats;
        let g = spark::Graph {
            min: Some(0.),
            max: Some(1.),
        };
        let graphs: Vec<String> = self
            .history
            .iter()
            .map(|history| {
                let sg = g.render(&history.iter().map(|v| *v).collect::<Vec<f32>>());
                history
                    .iter()
                    .zip(sg.chars())
                    .map(|(v, g)| {
                        let c = color_lerp(ColorRGB::green(), ColorRGB::red(), *v);
                        format!(
                            r##"<span foreground="{}">{}</span>"##,
                            color_to_string(c),
                            g
                        )
                    })
                    .collect()
            })
            .collect();
        let joiner = if self.history[0].len() > 1 { "|" } else { "" };
        Some(format!("CPU|{}", graphs.join(joiner)))
    }
}

impl Widget for CpuWidget {
    fn update(&mut self) -> Option<WidgetUpdate> {
        if let Some(history) = self.draw_history() {
            let mut data = Block::new();

            data.use_pango();
            data.append_full_text(&history);

            return Some(WidgetUpdate {
                refresh_interval: std::time::Duration::new(1, 0),
                data: Some(data),
            });
        }

        None
    }
}