harmony/harmony_cli/src/progress.rs

use indicatif::{MultiProgress, ProgressBar};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::Duration;

pub trait ProgressTracker: Send + Sync {
    fn contains_section(&self, id: &str) -> bool;
    fn add_section(&self, id: &str, message: &str);
    fn add_task(&self, section_id: &str, task_id: &str, message: &str);
    fn finish_task(&self, id: &str, message: &str);
    fn fail_task(&self, id: &str, message: &str);
    fn skip_task(&self, id: &str, message: &str);
    fn clear(&self);
}

struct Section {
    header_index: usize,
    task_count: usize,
    pb: ProgressBar,
}

struct IndicatifProgressTrackerState {
    sections: HashMap<String, Section>,
    tasks: HashMap<String, ProgressBar>,
    pb_count: usize,
}

#[derive(Clone)]
pub struct IndicatifProgressTracker {
    mp: MultiProgress,
    state: Arc<Mutex<IndicatifProgressTrackerState>>,
}

impl IndicatifProgressTracker {
    pub fn new(base: MultiProgress) -> Self {
        // The indicatif log bridge will insert a progress bar at the top.
        // To prevent our first section from being erased, we need to create
        // a dummy progress bar as our first progress bar.
        let _ = base.clear();
        let log_pb = base.add(ProgressBar::new(1));

        let mut sections = HashMap::new();
        sections.insert(
            "__log__".into(),
            Section {
                header_index: 0,
                task_count: 0,
                pb: log_pb.clone(),
            },
        );

        let mut tasks = HashMap::new();
        tasks.insert("__log__".into(), log_pb);

        let state = Arc::new(Mutex::new(IndicatifProgressTrackerState {
            sections,
            tasks,
            pb_count: 1,
        }));

        Self { mp: base, state }
    }
}

impl ProgressTracker for IndicatifProgressTracker {
    fn add_section(&self, id: &str, message: &str) {
        let mut state = self.state.lock().unwrap();

        let header_pb = self
            .mp
            .add(ProgressBar::new(1).with_style(crate::theme::SECTION_STYLE.clone()));
        header_pb.finish_with_message(message.to_string());

        let header_index = state.pb_count;
        state.pb_count += 1;

        state.sections.insert(
            id.to_string(),
            Section {
                header_index,
                task_count: 0,
                pb: header_pb,
            },
        );
    }

    fn add_task(&self, section_id: &str, task_id: &str, message: &str) {
        let mut state = self.state.lock().unwrap();

        let insertion_index = {
            let current_section = state
                .sections
                .get(section_id)
                .expect("Section ID not found");
            current_section.header_index + current_section.task_count + 1 // +1 to insert after header
        };

        let pb = self.mp.insert(insertion_index, ProgressBar::new_spinner());
        pb.set_style(crate::theme::SPINNER_STYLE.clone());
        pb.set_prefix("  ");
        pb.set_message(message.to_string());
        pb.enable_steady_tick(Duration::from_millis(80));

        state.pb_count += 1;

        let section = state
            .sections
            .get_mut(section_id)
            .expect("Section ID not found");
        section.task_count += 1;

        // We inserted a new progress bar, so we must update the header_index
        // for all subsequent sections.
        for (id, s) in state.sections.iter_mut() {
            if id != section_id && s.header_index >= insertion_index {
                s.header_index += 1;
            }
        }

        state.tasks.insert(task_id.to_string(), pb);
    }

    fn finish_task(&self, id: &str, message: &str) {
        let state = self.state.lock().unwrap();
        if let Some(pb) = state.tasks.get(id) {
            pb.set_style(crate::theme::SUCCESS_SPINNER_STYLE.clone());
            pb.finish_with_message(message.to_string());
        }
    }

    fn fail_task(&self, id: &str, message: &str) {
        let state = self.state.lock().unwrap();
        if let Some(pb) = state.tasks.get(id) {
            pb.set_style(crate::theme::ERROR_SPINNER_STYLE.clone());
            pb.finish_with_message(message.to_string());
        }
    }

    fn skip_task(&self, id: &str, message: &str) {
        let state = self.state.lock().unwrap();
        if let Some(pb) = state.tasks.get(id) {
            pb.set_style(crate::theme::SKIP_SPINNER_STYLE.clone());
            pb.finish_with_message(message.to_string());
        }
    }

    fn contains_section(&self, id: &str) -> bool {
        let state = self.state.lock().unwrap();
        state.sections.contains_key(id)
    }

    fn clear(&self) {
        let mut state = self.state.lock().unwrap();

        state.tasks.values().for_each(|p| self.mp.remove(p));
        state.tasks.clear();
        state.sections.values().for_each(|s| self.mp.remove(&s.pb));
        state.sections.clear();
        state.pb_count = 0;

        let _ = self.mp.clear();
    }
}