mimir_progress 0.1.0

// Copyright 2025 Red Hat, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::{
    io::{stdout, Write},
    time::{Duration, Instant},
};

pub struct Progress {
    start_time: Instant,
    end_time: Instant,
    last_tick_time: Instant,
    pub total: usize,
    current: usize,
    interval: ProgressInterval,
    prefix: String,
    autoprint: bool,
    finished: bool,
    same_line: bool,
}

/// Defines the "refresh rate" of the progress updates, by defining either a time span between
/// updates, a number of discrete updates required between updates, or a percentage of the total
/// number of items.
pub enum ProgressInterval {
    /// The Time interval causes an update to be printed every N units of time.
    Time(Duration),
    /// The Count interval causes an update to be printed when the count has been incremented
    /// N times.
    Count(usize),
    /// The Percent interval causes an update to be printed when the current count has
    /// incremented by N percent of the total count.
    Percent(usize),
}

impl ProgressInterval {
    /// Return true if an interval boundary has been reached.
    fn should_update(&self, progress: &Progress) -> bool {
        match self {
            ProgressInterval::Time(interval) => {
                let now = Instant::now();
                &now.duration_since(progress.last_tick_time) > interval
            }
            ProgressInterval::Count(interval) => progress.current % interval == 0,
            ProgressInterval::Percent(interval) => {
                let current_percent = (100 * progress.current) / progress.total;
                current_percent % interval == 0
            }
        }
    }
}

impl Progress {
    pub fn new(total: usize, interval: ProgressInterval) -> Progress {
        let now = Instant::now();

        Progress {
            start_time: now,
            end_time: now, // will be overwritten later
            total,
            current: 0,
            interval,
            prefix: "".to_string(),
            finished: false,
            last_tick_time: now,
            autoprint: true,
            same_line: true,
        }
    }
    pub fn set_prefix(&mut self, msg: String) {
        self.prefix = msg;
    }
    /// true causes Progress to print to stdout, false causes it to not print.
    pub fn set_autoprint(&mut self, enabled: bool) {
        self.autoprint = enabled;
    }

    /// Print updates on the same line in the terminal, as a typical progress bar would.  In some
    /// CI environments (like AWS CodeBuild) this should be disabled in order for progress updates
    /// to appear.
    pub fn set_same_line(&mut self, same_line: bool) {
        self.same_line = same_line;
    }

    /// Assign a value to the current amount.
    /// Returns `true` if a new line should be printed based on the chosen interval.
    pub fn set(&mut self, new_current: usize) -> bool {
        let mut should_print = false;

        if !self.finished {
            let prev = self.current;
            self.current = new_current;
            should_print = self.post_update(prev);
        }

        should_print
    }

    /// Increment the current count by 1.
    /// Returns `true` if a new line should be printed based on the chosen interval.
    pub fn inc(&mut self) -> bool {
        self.set(self.current + 1)
    }

    /// Handle tasks after the current count has been changed, such as updating the latest
    /// timestamp.  Returns `true` if the post_update step determines that a new item should be
    /// printed based on the chosen interval.
    fn post_update(&mut self, prev: usize) -> bool {
        let mut should_print = false;

        if self.should_update(prev) {
            self.last_tick_time = Instant::now();
            should_print = true;
            if self.autoprint {
                self.print();
            }
        }

        if self.current >= self.total {
            self.finished = true;
            self.end_time = Instant::now();
            if self.autoprint {
                println!("{}", self.finish_msg());
            }
        }
        should_print
    }

    /// Get the current count.
    pub fn current_count(&self) -> usize {
        self.current
    }

    /// Returns true if a progress update should be logged (either printed to the screen if
    /// `autoprint` is enabled, or returned by `msg()`).  A variety of conditions determine when
    /// updates are logged, including the total count being reached, or the user-selected interval
    /// reaching a boundary.  For example, if the user chose an interval of
    /// `Time(Duration::from_secs(2))` then an elapse of 2 seconds represents a boundary of that
    /// interval, and an update will be logged.
    fn should_update(&self, prev_count: usize) -> bool {
        let still_running = !self.finished;
        let just_finished = self.current >= self.total;
        let on_interval_boundary = self.interval.should_update(self);
        let just_started = prev_count == 0;

        still_running && (just_started || just_finished || on_interval_boundary)
    }

    fn print(&self) {
        if self.same_line {
            print!("{}\r{}", termion::clear::CurrentLine, self.msg());
            stdout().flush().expect("stdout flush failed");
        } else {
            println!("{}", self.msg());
        }
    }

    /// Get the current progress message.
    pub fn msg(&self) -> String {
        let complete_char = "\u{2588}";
        let incomplete_char = "\u{2591}";

        let percent = (100.0 * self.current as f64 / self.total as f64) as usize;

        let scaled_percent = percent / 4;

        let bar = format!(
            "{}{}",
            complete_char.repeat(scaled_percent),
            incomplete_char.repeat(25 - scaled_percent)
        );

        format!(
            "{msg}{current_pad}{current}/{total}: {bar} {percent_pad}{percent}% ({elapsed} elapsed, {eta} remaining, {hz:.2}/s)",
            msg = self.prefix,
            current_pad = if self.current > 0 {
                " ".repeat((self.total.ilog10() - self.current.ilog10()) as usize)
            } else { "".to_string() },
            current = self.current,
            total = self.total,
            percent = percent,
            percent_pad = if percent > 0 {
                " ".repeat(2 - percent.ilog10() as usize)
            } else { "".to_string() },
            eta = humantime::format_duration(self.time_estimate()),
            elapsed = humantime::format_duration(Duration::from_secs(
                self.start_time.elapsed().as_secs()
            )),
            hz = self.current as f32 / self.start_time.elapsed().as_secs_f32()
        )
    }

    pub fn finish_msg(&self) -> String {
        format!(
            "{}{msg}finished in {dur}",
            if self.same_line { "\n" } else { "" },
            msg = self.prefix,
            dur = humantime::format_duration(Duration::from_secs(
                self.end_time.duration_since(self.start_time).as_secs()
            )),
        )
    }

    /// Returns an estimate of the remaining time.
    fn time_estimate(&self) -> Duration {
        let remaining_count = self.total - self.current;
        let elapsed_secs = self.start_time.elapsed().as_secs_f64();

        Duration::from_secs((remaining_count as f64 * elapsed_secs / self.current as f64) as u64)
    }
}

#[cfg(test)]
mod tests {

    use std::thread;

    use super::*;

    #[test]
    fn over_increment() {
        // increment too many times
        let mut pb = Progress::new(4, ProgressInterval::Count(2));
        for _ in 0..100 {
            pb.inc();
        }
    }

    #[test]
    fn count_interval_test() {
        let mut pb = Progress::new(10, ProgressInterval::Count(2));
        let mut msgs = Vec::new();
        for _ in 0..10 {
            if pb.inc() {
                msgs.push(pb.msg());
            }
        }
        assert_eq!(msgs.len(), 6);
        // Disabled exact tests because the rate (hz) varies from run to run.
        // assert_eq!(
        //     msgs,
        //     vec![
        //         "1/10: 10% (0s remaining)".to_string(),
        //         "2/10: 20% (0s remaining)".to_string(),
        //         "4/10: 40% (0s remaining)".to_string(),
        //         "6/10: 60% (0s remaining)".to_string(),
        //         "8/10: 80% (0s remaining)".to_string(),
        //         "10/10: 100% (0s remaining)".to_string(),
        //     ]
        // );
    }

    #[test]
    fn percent_interval_test() {
        let mut pb = Progress::new(10, ProgressInterval::Percent(50));

        let mut msgs = Vec::new();

        for _ in 0..10 {
            if pb.inc() {
                msgs.push(pb.msg());
            }
        }

        assert_eq!(msgs.len(), 3);
        // Disabled exact tests because the rate (hz) varies from run to run.
        // assert_eq!(
        //     msgs,
        //     vec![
        //         "1/10: 10% (0s remaining)".to_string(),
        //         "5/10: 50% (0s remaining)".to_string(),
        //         "10/10: 100% (0s remaining)".to_string(),
        //     ]
        // );
    }

    #[test]
    fn time_interval_test() {
        let mut pb = Progress::new(8, ProgressInterval::Time(Duration::from_millis(50)));

        let mut msgs = Vec::new();

        for _ in 0..8 {
            thread::sleep(Duration::from_millis(30));
            if pb.inc() {
                msgs.push(pb.msg());
            }
        }

        assert_eq!(msgs.len(), 5);
        // Disabled exact tests because the rate (hz) varies from run to run.
        // assert_eq!(
        //     msgs,
        //     vec![
        //         "1/8: ███░░░░░░░░░░░░░░░░░░░░░░  12% (0s elapsed, 0s remaining, 33.23/s)".to_string(),
        //         "3/8: █████████░░░░░░░░░░░░░░░░  37% (0s elapsed, 0s remaining, 33.23/s)".to_string(),
        //         "5/8: ███████████████░░░░░░░░░░  62% (0s elapsed, 0s remaining, 33.14/s)".to_string(),
        //         "7/8: █████████████████████░░░░  87% (0s elapsed, 0s remaining, 33.13/s)".to_string(),
        //         "8/8: █████████████████████████ 100% (0s elapsed, 0s remaining, 33.11/s)".to_string(),
        //     ]
        // );
    }
}