abbaye/builders/

orchestrator.rs

use std::collections::HashMap;
use std::time::Duration;

use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
use miette::{IntoDiagnostic, Result};
use tokio::sync::{mpsc, watch};
use tokio::task::JoinSet;
use tracing::{info, warn};

use crate::builders::{ArtifactPath, BuilderEntry, LogEvent};
use crate::cli::{COLOURS, GREEN, RED, RESET, YELLOW};
use crate::site::spinner_style;

/// Run all builders in parallel, respecting `depends_on` ordering.
///
/// Returns `(dist_artifacts, doc_artifacts)` – file-type artifacts go into
/// the dist/ directory while directory-type artifacts go into docs/.
pub(crate) async fn run_builders(
    builders: &[BuilderEntry],
    version: &str,
) -> Result<(Vec<ArtifactPath>, Vec<ArtifactPath>)> {
    let id_to_idx: HashMap<&str, usize> = builders
        .iter()
        .enumerate()
        .filter_map(|(i, e)| e.id.as_deref().map(|id| (id, i)))
        .collect();

    for (i, entry) in builders.iter().enumerate() {
        for dep in &entry.depends_on {
            if !id_to_idx.contains_key(dep.as_str()) {
                return Err(miette::miette!(
                    "builder #{i} ({}) lists '{}' in depends_on, \
                     but no builder has that id",
                    entry.label(),
                    dep
                ));
            }
        }
    }

    {
        let n = builders.len();
        let mut state = vec![0u8; n];

        fn dfs(
            idx: usize,
            id_to_idx: &HashMap<&str, usize>,
            builders: &[BuilderEntry],
            state: &mut Vec<u8>,
        ) -> Result<()> {
            if state[idx] == 1 {
                return Err(miette::miette!(
                    "dependency cycle detected involving builder #{idx} ({})",
                    builders[idx].id.as_deref().unwrap_or(builders[idx].label())
                ));
            }
            if state[idx] == 2 {
                return Ok(());
            }
            state[idx] = 1;
            for dep in &builders[idx].depends_on {
                if let Some(&dep_idx) = id_to_idx.get(dep.as_str()) {
                    dfs(dep_idx, id_to_idx, builders, state)?;
                }
            }
            state[idx] = 2;
            Ok(())
        }

        for i in 0..n {
            dfs(i, &id_to_idx, builders, &mut state)?;
        }
    }

    let mut completion_txs: HashMap<String, watch::Sender<Option<bool>>> = HashMap::new();
    let mut completion_rxs: HashMap<String, watch::Receiver<Option<bool>>> = HashMap::new();

    for entry in builders {
        if let Some(id) = &entry.id {
            let (tx, rx) = watch::channel(None::<bool>);
            completion_txs.insert(id.clone(), tx);
            completion_rxs.insert(id.clone(), rx);
        }
    }

    let total = builders.len();

    if crate::utils::is_interactive() {
        run_with_bars(builders, version, completion_txs, completion_rxs, total).await
    } else {
        run_with_logs(builders, version, completion_txs, completion_rxs, total).await
    }
}

// ── Shared builder-task body ────────────────────────────────────────────────

/// Events reported by the shared builder loop to the UI layer.
enum BuildEvent {
    DepWait(String),
    Skip(String),
    Run,
    Success(usize),
    Failure(String),
}

/// Run a single builder inside a `JoinSet` task, handling dependency waiting,
/// completion signalling, and reporting events back to the UI layer.
async fn run_builder_task(
    entry: BuilderEntry,
    version: String,
    dep_receivers: Vec<(String, watch::Receiver<Option<bool>>)>,
    my_tx: Option<watch::Sender<Option<bool>>>,
    log_tx: mpsc::UnboundedSender<LogEvent>,
    on_event: impl Fn(BuildEvent) + Send + 'static,
) -> Result<Vec<ArtifactPath>> {
    for (dep_id, mut rx) in dep_receivers {
        on_event(BuildEvent::DepWait(dep_id.clone()));

        let resolved = rx.wait_for(|v| v.is_some()).await;

        let succeeded = match resolved {
            Err(_) => false,
            Ok(r) => r.unwrap_or(false),
        };

        if !succeeded {
            on_event(BuildEvent::Skip(dep_id.clone()));
            if let Some(tx) = &my_tx {
                let _ = tx.send(Some(false));
            }
            return Ok(vec![]);
        }
    }

    on_event(BuildEvent::Run);
    let result = entry.build(&version, log_tx).await;
    let succeeded = result.is_ok();

    if let Some(tx) = my_tx {
        let _ = tx.send(Some(succeeded));
    }

    match &result {
        Ok(artifacts) => on_event(BuildEvent::Success(artifacts.len())),
        Err(e) => on_event(BuildEvent::Failure(e.to_string())),
    }
    result
}

/// Collect finished `JoinSet` tasks and split results into dist / doc artifacts.
async fn collect_results(
    mut join_set: JoinSet<Result<Vec<ArtifactPath>>>,
    on_done: impl FnOnce(bool),
) -> Result<(Vec<ArtifactPath>, Vec<ArtifactPath>)> {
    let mut errors: Vec<miette::Report> = Vec::new();
    let mut dist_artifacts = Vec::new();
    let mut doc_artifacts = Vec::new();

    while let Some(res) = join_set.join_next().await {
        match res.into_diagnostic()? {
            Ok(artifacts) => {
                for artifact in artifacts {
                    if artifact.path.is_dir() {
                        doc_artifacts.push(artifact);
                    } else {
                        dist_artifacts.push(artifact);
                    }
                }
            }
            Err(e) => errors.push(e),
        }
    }

    on_done(errors.is_empty());

    if let Some(first_err) = errors.into_iter().next() {
        return Err(first_err);
    }

    Ok((dist_artifacts, doc_artifacts))
}

// ── Interactive (progress bar) UI ───────────────────────────────────────────

async fn run_with_bars(
    builders: &[BuilderEntry],
    version: &str,
    mut completion_txs: HashMap<String, watch::Sender<Option<bool>>>,
    completion_rxs: HashMap<String, watch::Receiver<Option<bool>>>,
    total: usize,
) -> Result<(Vec<ArtifactPath>, Vec<ArtifactPath>)> {
    let multi = MultiProgress::new();

    let summary = multi.add(ProgressBar::new(total as u64));
    summary.set_style(
        ProgressStyle::with_template("{pos}/{len} builders  {bar:20.green/white}  {msg}")
            .expect("valid template"),
    );
    summary.set_message("building…");

    let mut join_set: JoinSet<Result<Vec<ArtifactPath>>> = JoinSet::new();

    for (i, entry) in builders.iter().enumerate() {
        let color = COLOURS[i % COLOURS.len()];
        let label = entry.id.as_deref().unwrap_or(entry.label());
        let colored_prefix = format!("{color}[{label}]{RESET}");

        let pb = multi.insert_before(&summary, ProgressBar::new_spinner());
        pb.set_style(spinner_style(false));
        pb.set_prefix(colored_prefix);
        pb.set_message("starting…");
        pb.enable_steady_tick(Duration::from_millis(100));

        let (log_tx, mut log_rx) = mpsc::unbounded_channel::<LogEvent>();

        let pb_log = pb.clone();
        let multi_log = multi.clone();
        let parent_color_idx = i;
        tokio::spawn(async move {
            let mut child_pbs: HashMap<String, ProgressBar> = HashMap::new();
            let mut last_child_pb = pb_log.clone();
            let mut child_color_idx = parent_color_idx + 1;

            while let Some(event) = log_rx.recv().await {
                match event {
                    LogEvent::Line(line) => {
                        pb_log.set_message(line);
                    }
                    LogEvent::ChildStart { id, label } => {
                        let child_color = COLOURS[child_color_idx % COLOURS.len()];
                        child_color_idx += 1;
                        let child_pb =
                            multi_log.insert_after(&last_child_pb, ProgressBar::new_spinner());
                        child_pb.set_style(spinner_style(true));
                        child_pb.set_prefix(format!("{child_color}[{label}]{RESET}"));
                        child_pb.set_message("starting…");
                        child_pb.enable_steady_tick(Duration::from_millis(100));
                        last_child_pb = child_pb.clone();
                        child_pbs.insert(id, child_pb);
                    }
                    LogEvent::ChildLine { id, line } => {
                        if let Some(child_pb) = child_pbs.get(&id) {
                            child_pb.set_message(line);
                        }
                    }
                    LogEvent::ChildFinish {
                        id,
                        success,
                        summary,
                    } => {
                        if let Some(child_pb) = child_pbs.remove(&id) {
                            if success {
                                child_pb.finish_with_message(format!(
                                    "{GREEN}\u{2713}{RESET} {summary}"
                                ));
                            } else {
                                child_pb
                                    .finish_with_message(format!("{RED}\u{2717}{RESET} {summary}"));
                            }
                        }
                    }
                }
            }
        });

        let dep_receivers: Vec<(String, watch::Receiver<Option<bool>>)> = entry
            .depends_on
            .iter()
            .filter_map(|dep_id| {
                completion_rxs
                    .get(dep_id)
                    .map(|rx| (dep_id.clone(), rx.clone()))
            })
            .collect();

        let my_tx: Option<watch::Sender<Option<bool>>> =
            entry.id.as_ref().and_then(|id| completion_txs.remove(id));

        let entry = entry.clone();
        let version = version.to_owned();
        let pb_task = pb.clone();
        let summary_task = summary.clone();

        join_set.spawn(async move {
            run_builder_task(
                entry,
                version,
                dep_receivers,
                my_tx,
                log_tx,
                move |event| match event {
                    BuildEvent::DepWait(dep_id) => {
                        pb_task.set_message(format!("waiting for '{dep_id}'…"));
                    }
                    BuildEvent::Skip(dep_id) => {
                        summary_task.inc(1);
                        pb_task.finish_with_message(format!(
                            "{YELLOW}\u{29B8} skipped{RESET} (dependency '{dep_id}' failed)"
                        ));
                    }
                    BuildEvent::Run => {
                        pb_task.set_message("running…");
                    }
                    BuildEvent::Success(count) => {
                        summary_task.inc(1);
                        pb_task.finish_with_message(format!(
                            "{GREEN}\u{2713} done{RESET} ({count} artifact(s))"
                        ));
                    }
                    BuildEvent::Failure(error) => {
                        summary_task.inc(1);
                        pb_task
                            .finish_with_message(format!("{RED}\u{2717} failed:{RESET} {error}"));
                    }
                },
            )
            .await
        });
    }

    let summary_bar = summary.clone();
    collect_results(join_set, move |ok| {
        let msg = if ok {
            format!("{GREEN}\u{2713} all done{RESET}")
        } else {
            format!("{RED}\u{2717} some builders failed{RESET}")
        };
        summary_bar.finish_with_message(msg);
    })
    .await
}

// ── Non-interactive (tracing) UI ────────────────────────────────────────────

async fn run_with_logs(
    builders: &[BuilderEntry],
    version: &str,
    mut completion_txs: HashMap<String, watch::Sender<Option<bool>>>,
    completion_rxs: HashMap<String, watch::Receiver<Option<bool>>>,
    total: usize,
) -> Result<(Vec<ArtifactPath>, Vec<ArtifactPath>)> {
    info!("Running {total} builder(s) …");

    let mut join_set: JoinSet<Result<Vec<ArtifactPath>>> = JoinSet::new();

    for entry in builders.iter() {
        let label = entry.id.as_deref().unwrap_or(entry.label()).to_owned();

        let (log_tx, mut log_rx) = mpsc::unbounded_channel::<LogEvent>();

        let log_label = label.clone();
        tokio::spawn(async move {
            while let Some(event) = log_rx.recv().await {
                match event {
                    LogEvent::Line(line) => info!("[{log_label}] {line}"),
                    LogEvent::ChildStart { label: child, .. } => {
                        info!("[{log_label}] [{child}] starting …");
                    }
                    LogEvent::ChildLine { id, line } => {
                        info!("[{log_label}] [{id}] {line}");
                    }
                    LogEvent::ChildFinish {
                        id,
                        success,
                        summary,
                    } => {
                        if success {
                            info!("[{log_label}] [{id}] {summary}");
                        } else {
                            warn!("[{log_label}] [{id}] failed: {summary}");
                        }
                    }
                }
            }
        });

        let dep_receivers: Vec<(String, watch::Receiver<Option<bool>>)> = entry
            .depends_on
            .iter()
            .filter_map(|dep_id| {
                completion_rxs
                    .get(dep_id)
                    .map(|rx| (dep_id.clone(), rx.clone()))
            })
            .collect();

        let my_tx: Option<watch::Sender<Option<bool>>> =
            entry.id.as_ref().and_then(|id| completion_txs.remove(id));

        let entry = entry.clone();
        let version = version.to_owned();
        let task_label = label.clone();

        join_set.spawn(async move {
            run_builder_task(
                entry,
                version,
                dep_receivers,
                my_tx,
                log_tx,
                move |event| match event {
                    BuildEvent::DepWait(dep_id) => {
                        info!("[{task_label}] waiting for '{dep_id}' …");
                    }
                    BuildEvent::Skip(dep_id) => {
                        info!("[{task_label}] skipped (dependency '{dep_id}' failed)");
                    }
                    BuildEvent::Run => {
                        info!("[{task_label}] running …");
                    }
                    BuildEvent::Success(count) => {
                        info!("[{task_label}] done ({count} artifact(s))");
                    }
                    BuildEvent::Failure(error) => {
                        warn!("[{task_label}] failed: {error}");
                    }
                },
            )
            .await
        });
    }

    collect_results(join_set, |ok| {
        if ok {
            info!("All builders completed successfully.");
        } else {
            warn!("Some builders failed.");
        }
    })
    .await
}