canic-cli 0.36.12

use canic_host::registry::RegistryEntry;
use std::{
    collections::{BTreeMap, BTreeSet},
    error::Error,
    fmt,
};

const TREE_BRANCH: &str = "├─ ";
const TREE_LAST: &str = "└─ ";
const TREE_PIPE: &str = "│  ";
const TREE_SPACE: &str = "   ";

///
/// RegistryTreeError
///

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum RegistryTreeError {
    CanisterNotInRegistry(String),
    AmbiguousRole {
        role: String,
        canisters: Vec<String>,
    },
}

impl fmt::Display for RegistryTreeError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::CanisterNotInRegistry(canister) => {
                write!(
                    f,
                    "registry JSON did not contain the requested canister or role {canister}"
                )
            }
            Self::AmbiguousRole { role, canisters } => write!(
                f,
                "registry role {role} has multiple canisters; pass one canister principal: {}",
                canisters.join(", ")
            ),
        }
    }
}

impl Error for RegistryTreeError {}

///
/// RegistryRow
///

pub struct RegistryRow<'a> {
    pub entry: &'a RegistryEntry,
    pub tree_prefix: String,
}

pub fn registry_rows(registry: &[RegistryEntry]) -> Vec<RegistryRow<'_>> {
    let roots = root_entries(registry);
    let children = child_entries(registry);
    let mut entries = Vec::new();

    for root in roots {
        collect_visible_entry(root, &children, "", "", &mut entries);
    }

    entries
}

pub fn visible_entries<'a>(
    registry: &'a [RegistryEntry],
    canister: Option<&str>,
) -> Result<Vec<&'a RegistryEntry>, RegistryTreeError> {
    Ok(visible_rows(registry, canister)?
        .into_iter()
        .map(|row| row.entry)
        .collect())
}

pub fn visible_rows<'a>(
    registry: &'a [RegistryEntry],
    anchor: Option<&str>,
) -> Result<Vec<RegistryRow<'a>>, RegistryTreeError> {
    let Some(anchor) = anchor else {
        return Ok(registry_rows(registry));
    };

    let by_pid = registry
        .iter()
        .map(|entry| (entry.pid.as_str(), entry))
        .collect::<BTreeMap<_, _>>();
    let root = by_pid
        .get(anchor)
        .copied()
        .map_or_else(|| unique_role_entry(registry, anchor), Ok)?;
    let children = child_entries(registry);
    let mut entries = Vec::new();
    collect_visible_entry(root, &children, "", "", &mut entries);
    Ok(entries)
}

fn unique_role_entry<'a>(
    registry: &'a [RegistryEntry],
    role: &str,
) -> Result<&'a RegistryEntry, RegistryTreeError> {
    let matches = registry
        .iter()
        .filter(|entry| entry.role.as_deref() == Some(role))
        .collect::<Vec<_>>();
    match matches.as_slice() {
        [entry] => Ok(*entry),
        [] => Err(RegistryTreeError::CanisterNotInRegistry(role.to_string())),
        entries => Err(RegistryTreeError::AmbiguousRole {
            role: role.to_string(),
            canisters: entries.iter().map(|entry| entry.pid.clone()).collect(),
        }),
    }
}

fn root_entries(registry: &[RegistryEntry]) -> Vec<&RegistryEntry> {
    let ids = registry
        .iter()
        .map(|entry| entry.pid.as_str())
        .collect::<BTreeSet<_>>();
    registry
        .iter()
        .filter(|entry| {
            entry
                .parent_pid
                .as_deref()
                .is_none_or(|parent| !ids.contains(parent))
        })
        .collect()
}

fn child_entries(registry: &[RegistryEntry]) -> BTreeMap<&str, Vec<&RegistryEntry>> {
    let mut children = BTreeMap::<&str, Vec<&RegistryEntry>>::new();
    for entry in registry {
        if let Some(parent) = entry.parent_pid.as_deref() {
            children.entry(parent).or_default().push(entry);
        }
    }
    for entries in children.values_mut() {
        entries.sort_by_key(|entry| (entry.role.as_deref().unwrap_or(""), entry.pid.as_str()));
    }
    children
}

fn collect_visible_entry<'a>(
    entry: &'a RegistryEntry,
    children: &BTreeMap<&str, Vec<&'a RegistryEntry>>,
    tree_prefix: &str,
    child_prefix: &str,
    entries: &mut Vec<RegistryRow<'a>>,
) {
    entries.push(RegistryRow {
        entry,
        tree_prefix: tree_prefix.to_string(),
    });
    if let Some(child_entries) = children.get(entry.pid.as_str()) {
        for (index, child) in child_entries.iter().enumerate() {
            let is_last = index + 1 == child_entries.len();
            let branch = if is_last { TREE_LAST } else { TREE_BRANCH };
            let carry = if is_last { TREE_SPACE } else { TREE_PIPE };
            let child_tree_prefix = format!("{child_prefix}{branch}");
            let descendant_prefix = format!("{child_prefix}{carry}");
            collect_visible_entry(
                child,
                children,
                &child_tree_prefix,
                &descendant_prefix,
                entries,
            );
        }
    }
}