harmony/harmony_cli/src/cli_logger.rs

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

use crate::progress::{IndicatifProgressTracker, ProgressTracker};

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

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

    handle
}

fn configure_logger() -> MultiProgress {
    let logger =
        env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).build();
    let level = logger.filter();
    let progress = MultiProgress::new();

    LogWrapper::new(progress.clone(), logger)
        .try_init()
        .unwrap();
    log::set_max_level(level);

    progress
}

async fn handle_events(base_progress: MultiProgress) {
    let progress_tracker = Arc::new(IndicatifProgressTracker::new(base_progress.clone()));
    let preparing_topology = Arc::new(Mutex::new(false));
    let current_score: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));

    instrumentation::subscribe("Harmony CLI Logger", {
        move |event| {
            let progress_tracker = Arc::clone(&progress_tracker);
            let preparing_topology = Arc::clone(&preparing_topology);
            let current_score = Arc::clone(&current_score);

            async move {
                let mut preparing_topology = preparing_topology.lock().unwrap();
                let mut current_score = current_score.lock().unwrap();

                match event {
                    HarmonyEvent::HarmonyStarted => {}
                    HarmonyEvent::HarmonyFinished => {
                        progress_tracker.add_section(
                            "harmony-summary",
                            &format!("\n{} Harmony completed\n\n", crate::theme::EMOJI_HARMONY),
                        );
                        progress_tracker.add_section("harmony-finished", "\n\n");
                        return false;
                    }
                    HarmonyEvent::TopologyStateChanged {
                        topology,
                        status,
                        message,
                    } => {
                        let section_key = topology_key(&topology);

                        match status {
                            TopologyStatus::Queued => {}
                            TopologyStatus::Preparing => {
                                progress_tracker.add_section(
                                    &section_key,
                                    &format!(
                                        "\n{} Preparing environment: {topology}...",
                                        crate::theme::EMOJI_TOPOLOGY
                                    ),
                                );
                                (*preparing_topology) = true;
                            }
                            TopologyStatus::Success => {
                                (*preparing_topology) = false;
                                progress_tracker.add_task(&section_key, "topology-success", "");
                                progress_tracker
                                    .finish_task("topology-success", &message.unwrap_or("".into()));
                            }
                            TopologyStatus::Noop => {
                                (*preparing_topology) = false;
                                progress_tracker.add_task(&section_key, "topology-skip", "");
                                progress_tracker
                                    .skip_task("topology-skip", &message.unwrap_or("".into()));
                            }
                            TopologyStatus::Error => {
                                progress_tracker.add_task(&section_key, "topology-error", "");
                                (*preparing_topology) = false;
                                progress_tracker
                                    .fail_task("topology-error", &message.unwrap_or("".into()));
                            }
                        }
                    }
                    HarmonyEvent::InterpretExecutionStarted {
                        execution_id: task_key,
                        topology,
                        interpret: _,
                        score,
                        message,
                    } => {
                        let is_key_topology = (*preparing_topology)
                            && progress_tracker.contains_section(&topology_key(&topology));
                        let is_key_current_score = current_score.is_some()
                            && progress_tracker
                                .contains_section(&score_key(&current_score.clone().unwrap()));
                        let is_key_score = progress_tracker.contains_section(&score_key(&score));

                        let section_key = if is_key_topology {
                            topology_key(&topology)
                        } else if is_key_current_score {
                            score_key(&current_score.clone().unwrap())
                        } else if is_key_score {
                            score_key(&score)
                        } else {
                            (*current_score) = Some(score.clone());
                            let key = score_key(&score);
                            progress_tracker.add_section(
                                &key,
                                &format!(
                                    "{} Interpreting score: {score}...",
                                    crate::theme::EMOJI_SCORE
                                ),
                            );
                            key
                        };

                        progress_tracker.add_task(&section_key, &task_key, &message);
                    }
                    HarmonyEvent::InterpretExecutionFinished {
                        execution_id: task_key,
                        topology: _,
                        interpret: _,
                        score,
                        outcome,
                    } => {
                        if current_score.is_some() && current_score.clone().unwrap() == score {
                            (*current_score) = None;
                        }

                        match outcome {
                            Ok(outcome) => match outcome.status {
                                harmony::interpret::InterpretStatus::SUCCESS => {
                                    progress_tracker.finish_task(&task_key, &outcome.message);
                                }
                                harmony::interpret::InterpretStatus::NOOP => {
                                    progress_tracker.skip_task(&task_key, &outcome.message);
                                }
                                _ => progress_tracker.fail_task(&task_key, &outcome.message),
                            },
                            Err(err) => {
                                progress_tracker.fail_task(&task_key, &err.to_string());
                            }
                        }
                    }
                }
                true
            }
        }
    })
    .await;
}

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

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