harmony/harmony_cli/src/cli_logger.rs

use harmony::{
    instrumentation::{self, HarmonyEvent},
    topology::TopologyStatus,
};
use indicatif::{MultiProgress, ProgressBar};
use indicatif_log_bridge::LogWrapper;
use std::{
    collections::HashMap,
    sync::{Arc, Mutex},
};

use crate::progress;

pub fn init() -> tokio::task::JoinHandle<()> {
    configure_logger();
    let handle = tokio::spawn(handle_events());

    loop {
        if instrumentation::instrument(HarmonyEvent::HarmonyStarted).is_ok() {
            break;
        }
    }

    handle
}

fn configure_logger() {
    let logger =
        env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).build();
    let level = logger.filter();
    let multi = MultiProgress::new();
    LogWrapper::new(multi.clone(), logger).try_init().unwrap();
    log::set_max_level(level);
}

async fn handle_events() {
    instrumentation::subscribe("Harmony CLI Logger", {
        let sections: Arc<Mutex<HashMap<String, MultiProgress>>> =
            Arc::new(Mutex::new(HashMap::new()));
        let progress_bars: Arc<Mutex<HashMap<String, ProgressBar>>> =
            Arc::new(Mutex::new(HashMap::new()));

        move |event| {
            let sections_clone = Arc::clone(&sections);
            let progress_bars_clone = Arc::clone(&progress_bars);

            async move {
                let mut sections = sections_clone.lock().unwrap();
                let mut progress_bars = progress_bars_clone.lock().unwrap();

                match event {
                    HarmonyEvent::HarmonyStarted => {}
                    HarmonyEvent::TopologyStateChanged {
                        topology,
                        status,
                        message,
                    } => {
                        let section_key = topology_key(&topology);

                        match status {
                            TopologyStatus::Queued => {}
                            TopologyStatus::Preparing => {
                                let section = progress::new_section(format!(
                                    "{} Preparing environment: {topology}...",
                                    crate::theme::EMOJI_TOPOLOGY,
                                ));
                                (*sections).insert(section_key, section);
                            }
                            TopologyStatus::Success => {
                                let section = (*sections).get(&section_key).unwrap();
                                let progress = progress::add_spinner(section, "".into());

                                progress::success(
                                    section,
                                    Some(progress),
                                    message.unwrap_or("".into()),
                                );

                                (*sections).remove(&section_key);
                            }
                            TopologyStatus::Noop => {
                                let section = (*sections).get(&section_key).unwrap();
                                let progress = progress::add_spinner(section, "".into());

                                progress::skip(
                                    section,
                                    Some(progress),
                                    message.unwrap_or("".into()),
                                );

                                (*sections).remove(&section_key);
                            }
                            TopologyStatus::Error => {
                                let section = (*sections).get(&section_key).unwrap();
                                let progress = progress::add_spinner(section, "".into());

                                progress::error(
                                    section,
                                    Some(progress),
                                    message.unwrap_or("".into()),
                                );

                                (*sections).remove(&section_key);
                            }
                        }
                    }
                    HarmonyEvent::InterpretExecutionStarted {
                        topology,
                        interpret,
                        score,
                        message,
                    } => {
                        let section_key = if (*sections).contains_key(&topology_key(&topology)) {
                            topology_key(&topology)
                        } else if (*sections).contains_key(&score_key(&score)) {
                            score_key(&interpret)
                        } else {
                            let key = score_key(&score);
                            let section = progress::new_section(format!(
                                "\n{} Interpreting score: {score}...",
                                crate::theme::EMOJI_SCORE,
                            ));
                            (*sections).insert(key.clone(), section);
                            key
                        };
                        let section = (*sections).get(&section_key).unwrap();
                        let progress_bar = progress::add_spinner(section, message);

                        (*progress_bars).insert(interpret_key(&interpret), progress_bar);
                    }
                    HarmonyEvent::InterpretExecutionFinished {
                        topology,
                        interpret,
                        score,
                        outcome,
                    } => {
                        let has_topology = (*sections).contains_key(&topology_key(&topology));
                        let section_key = if has_topology {
                            topology_key(&topology)
                        } else {
                            score_key(&score)
                        };

                        let section = (*sections).get(&section_key).unwrap();
                        let progress_bar =
                            (*progress_bars).get(&interpret_key(&interpret)).cloned();

                        let _ = section.clear();

                        match outcome {
                            Ok(outcome) => match outcome.status {
                                harmony::interpret::InterpretStatus::SUCCESS => {
                                    progress::success(section, progress_bar, outcome.message)
                                }
                                harmony::interpret::InterpretStatus::NOOP => {
                                    progress::skip(section, progress_bar, outcome.message)
                                }
                                _ => progress::error(section, progress_bar, outcome.message),
                            },
                            Err(err) => {
                                progress::error(section, progress_bar, err.to_string());
                            }
                        }

                        if !has_topology {
                            (*progress_bars).remove(&section_key);
                        }
                    }
                }
                true
            }
        }
    })
    .await;
}

fn topology_key(topology: &str) -> String {
    format!("topology-{topology}")
}

fn score_key(score: &str) -> String {
    format!("score-{score}")
}

fn interpret_key(interpret: &str) -> String {
    format!("interpret-{interpret}")
}