//! The graph itself, from the members alone: a ring by hash, plus each member's
//! lowest-scoring pairs, stable members preferred; per transport, and across them. Pure:
//! the same members give every node the same graph.

use sha2::{Digest, Sha256};

use super::*;

pub(super) const RING_LABEL: &[u8] = b"continuum/topology-ring/1";

pub(super) const PAIR_LABEL: &[u8] = b"continuum/topology-pair/2";

/// How many members may prefer the same stable member, as a multiple of `k`. Without a
/// bound, 20 stable members among 2000 get about 500 neighbours each: every
/// announcement in the swarm lands on them, and the backbone saturates first. Members
/// that find a stable member full prefer the next one, then go by pair score, so the
/// graph stays connected and symmetric.
pub(super) const STABLE_FANIN_FACTOR: usize = 4;

/// The transports `addresses` are on: their schemes (`mycelium` of `mycelium:<key>`);
/// the empty one for an address without (a transport whose identities are addresses).
pub(crate) fn schemes(addresses: &[String]) -> BTreeSet<String> {
    addresses
        .iter()
        .map(|address| {
            address
                .split_once(':')
                .map_or("", |(scheme, _)| scheme)
                .to_owned()
        })
        .collect()
}

/// The neighbours of `local` among `members` (by the transports each is on), on each of
/// `local_transports`: on each, its neighbours among the members on it
/// (`neighbours`), all together.
pub(crate) fn neighbours_across(
    members: &BTreeMap<String, BTreeSet<String>>,
    local: &str,
    local_transports: &BTreeSet<String>,
    k: usize,
    unreachable: &dyn Fn(&str) -> bool,
    stable: &dyn Fn(&str) -> bool,
) -> BTreeSet<String> {
    let mut all = BTreeSet::new();
    for transport in local_transports {
        let on_it: Vec<String> = members
            .iter()
            .filter(|(member, transports)| *member != local && transports.contains(transport))
            .map(|(member, _)| member.clone())
            .collect();
        all.extend(neighbours(&on_it, local, k, unreachable, stable));
    }
    all
}

/// The neighbours of `local` among `members` (which may or may not contain `local`),
/// skipping members for which `unreachable` holds.
pub(crate) fn neighbours(
    members: &[String],
    local: &str,
    k: usize,
    unreachable: &dyn Fn(&str) -> bool,
    stable: &dyn Fn(&str) -> bool,
) -> BTreeSet<String> {
    let mut all: Vec<String> = members.to_vec();
    all.push(local.to_owned());
    graph(&all, k, &|m: &str| m != local && unreachable(m), stable)
        .remove(local)
        .unwrap_or_default()
}

/// Every reachable member's neighbours. Hashes each member once and scores pairs by
/// mixing the two hashes (`pair_score`), so the N² scores are a few nanoseconds each:
/// tens of milliseconds for thousands of members. `stable` members (`swarm::stability`)
/// are preferred for half of each member's scored slots, each by at most
/// `STABLE_FANIN_FACTOR * k` members.
///
/// Each reachable member walks the ring (by key hash) both ways to the first reachable
/// member, which gives a ring over the reachable members, and adds the `k - 2`
/// reachable members with the lowest pair score. A link exists if either end chose it.
/// With at most `k + 1` reachable members, the reachable members form a full mesh.
///
/// Each node judges reachability by itself, so nodes may briefly disagree; that only
/// adds links (whoever chose a link initiates, the other answers), it never cuts one.
/// The preference budget is spent in key order, which every node shares, so nodes that
/// agree on reachability compute the same links.
pub(crate) fn graph(
    members: &[String],
    k: usize,
    unreachable: &dyn Fn(&str) -> bool,
    stable: &dyn Fn(&str) -> bool,
) -> BTreeMap<String, BTreeSet<String>> {
    let k = k.max(2);
    let all: BTreeSet<&str> = members.iter().map(String::as_str).collect();
    let live: BTreeSet<&str> = all.iter().copied().filter(|m| !unreachable(m)).collect();
    let mut links: BTreeMap<String, BTreeSet<String>> = live
        .iter()
        .map(|m| ((*m).to_owned(), BTreeSet::new()))
        .collect();

    let mut link = |a: &str, b: &str| {
        if a != b {
            links.get_mut(a).expect("live member").insert(b.to_owned());
            links.get_mut(b).expect("live member").insert(a.to_owned());
        }
    };

    if live.len() <= k + 1 {
        for a in &live {
            for b in &live {
                link(a, b);
            }
        }
        return links;
    }

    // The ring over all members, so its order doesn't change when one goes offline.
    let mut ring: Vec<([u8; 32], &str)> =
        all.iter().map(|m| (hash(RING_LABEL, &[*m]), *m)).collect();
    ring.sort_unstable();
    let n = ring.len();
    let position: BTreeMap<&str, usize> =
        ring.iter().enumerate().map(|(i, (_, m))| (*m, i)).collect();
    let first_live = |from: usize, step: usize| {
        (1..n)
            .map(|d| ring[(from + d * step) % n].1)
            .find(|m| live.contains(m))
            .expect("more than k + 1 live members")
    };

    // Live members in key order, with their pair keys; and those of them that are stable.
    let keyed: Vec<(&str, u64)> = live.iter().map(|m| (*m, pair_key(m))).collect();
    let stable_keyed: Vec<(&str, u64)> = keyed.iter().copied().filter(|(m, _)| stable(m)).collect();

    // How many members have preferred each stable member so far (`keyed` is in key
    // order, the same on every node).
    let budget = STABLE_FANIN_FACTOR * k;
    let mut preferred_by: BTreeMap<&str, usize> = BTreeMap::new();
    // Half the slots (rounded up) prefer stable members, so long-running nodes form a
    // backbone; the rest go by pair score alone, keeping paths short. A stable member
    // enough others prefer already is skipped, so no hub grows with the swarm.
    let stable_slots = (k - 2).div_ceil(2);

    for &(member, key) in &keyed {
        let at = position[member];
        let next = first_live(at, 1);
        let prev = first_live(at, n - 1);
        link(member, next);
        link(member, prev);
        let skip = [member, next, prev];
        // The stable candidates in score order (there are few).
        let mut stable_candidates = scored(&stable_keyed, key, skip);
        stable_candidates.sort_unstable();
        let preferred: Vec<&str> = stable_candidates
            .iter()
            .map(|(_, other)| *other)
            .filter(|other| preferred_by.get(other).copied().unwrap_or(0) < budget)
            .take(stable_slots)
            .collect();
        for other in preferred.iter().copied() {
            *preferred_by.entry(other).or_default() += 1;
        }
        // The rest: the lowest scores among the others that aren't preferred, which are
        // among the `need + preferred` lowest overall. Selecting those is linear; only
        // they are sorted.
        let need = k - 2 - preferred.len();
        let mut candidates = scored(&keyed, key, skip);
        let lowest = (need + preferred.len()).min(candidates.len());
        if lowest < candidates.len() {
            candidates.select_nth_unstable(lowest);
            candidates.truncate(lowest);
        }
        candidates.sort_unstable();
        let rest = candidates
            .iter()
            .map(|(_, other)| *other)
            .filter(|other| !preferred.contains(other))
            .take(need);
        for partner in preferred.iter().copied().chain(rest) {
            link(member, partner);
        }
    }
    links
}

/// `others` but those in `skip`, scored as pairs with the member whose key is `key`.
pub(super) fn scored<'a>(
    others: &[(&'a str, u64)],
    key: u64,
    skip: [&str; 3],
) -> Vec<(u64, &'a str)> {
    others
        .iter()
        .filter(|(other, _)| !skip.contains(other))
        .map(|&(other, other_key)| (mix_pair(key, other_key), other))
        .collect()
}

/// Symmetric: the same for (a, b) and (b, a).
pub(crate) fn pair_score(a: &str, b: &str) -> u64 {
    mix_pair(pair_key(a), pair_key(b))
}

/// What a member contributes to its pair scores: computed once per member.
pub(super) fn pair_key(member: &str) -> u64 {
    let digest = hash(PAIR_LABEL, &[member]);
    u64::from_be_bytes(digest[..8].try_into().expect("8 bytes"))
}

/// The score of a pair from its members' keys, symmetric: SplitMix64's finalizer over
/// the lower key and the mixed higher one. Keys are uniform (SHA-256), so for each member
/// the others' scores are in an order unrelated to anyone else's.
pub(super) fn mix_pair(a: u64, b: u64) -> u64 {
    let (low, high) = if a <= b { (a, b) } else { (b, a) };
    mix(low ^ mix(high))
}

pub(super) fn mix(mut z: u64) -> u64 {
    z = z.wrapping_add(0x9e37_79b9_7f4a_7c15);
    z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
    z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
    z ^ (z >> 31)
}

pub(super) fn hash(label: &[u8], fields: &[&str]) -> [u8; 32] {
    let mut hasher = Sha256::new();
    for field in [label]
        .into_iter()
        .chain(fields.iter().map(|f| f.as_bytes()))
    {
        hasher.update((field.len() as u32).to_be_bytes());
        hasher.update(field);
    }
    hasher.finalize().into()
}