//! Which members a node syncs with: its neighbours.
//!
//! In a full mesh every change reaches every node from every other node (Automerge
//! decides per pair, so nobody knows that a third peer already passed it on), so the
//! copies of one write grow with N². Above `k` members, each node keeps only a few
//! neighbours, and changes travel over several hops.
//!
//! Every node computes the graph from the member list and which members it can reach,
//! with no coordination:
//! * a ring ordered by key hash: each member's nearest reachable member in both
//!   directions, which keeps the reachable members connected;
//! * plus, for each member, the `k - 2` reachable others with the lowest pair score (a
//!   hash of both keys, so both ends agree on it), which keeps paths short.
//!
//! A link exists if either end chose it, so links are symmetric and a node has at least
//! `k` neighbours. With at most `k + 1` reachable members the graph is the full mesh:
//! small swarms behave exactly as before. A member joining, leaving or going offline
//! changes only links near it.

use std::collections::{BTreeMap, BTreeSet};

use sha2::{Digest, Sha256};
use tokio::time::Instant;
use tracing::info;

use crate::state::AppState;
use crate::swarm::members;

/// This node's current neighbours, and the member list they were computed from.
#[derive(Debug, Default)]
pub(crate) struct NeighbourView {
    members: Vec<String>,
    k: usize,
    unreachable: BTreeSet<String>,
    stable: BTreeSet<String>,
    neighbours: BTreeSet<String>,
    /// Every reachable member is a neighbour.
    full_mesh: bool,
}

/// Whether every reachable member is our neighbour (also while the view is unknown).
pub(crate) fn is_full_mesh(state: &AppState) -> bool {
    state
        .neighbours
        .read()
        .unwrap_or_else(|poisoned| poisoned.into_inner())
        .as_ref()
        .is_none_or(|view| view.full_mesh)
}

/// The current view in numbers, once known.
#[derive(Debug, Clone, Copy)]
pub(crate) struct Summary {
    /// Registered members, us included.
    pub(crate) members: usize,
    /// Members other than us not currently judged unreachable.
    pub(crate) reachable: usize,
    pub(crate) stable: usize,
    pub(crate) neighbours: usize,
    pub(crate) stable_neighbours: usize,
    pub(crate) full_mesh: bool,
}

pub(crate) fn summary(state: &AppState) -> Option<Summary> {
    let local = state.local_id();
    state
        .neighbours
        .read()
        .unwrap_or_else(|poisoned| poisoned.into_inner())
        .as_ref()
        .map(|view| Summary {
            members: view.members.len(),
            reachable: view
                .members
                .iter()
                .filter(|m| m.as_str() != local && !view.unreachable.contains(*m))
                .count(),
            stable: view.stable.len(),
            neighbours: view.neighbours.len(),
            stable_neighbours: view
                .neighbours
                .iter()
                .filter(|m| view.stable.contains(*m))
                .count(),
            full_mesh: view.full_mesh,
        })
}

/// Whether we initiate contact with `peer`: HELLOs, keepalives, dormancy checks, sync
/// openings. Everyone is a neighbour until the view is known (while joining, or before
/// any member registered), so joining works as in a full mesh. We always *answer*
/// anyone who contacts us.
pub(crate) fn is_neighbour(state: &AppState, peer: &str) -> bool {
    match &*state
        .neighbours
        .read()
        .unwrap_or_else(|poisoned| poisoned.into_inner())
    {
        Some(view) => view.neighbours.contains(peer),
        None => true,
    }
}

/// How the neighbours changed.
pub(crate) struct NeighbourChange {
    pub(crate) added: Vec<String>,
    pub(crate) removed: Vec<String>,
}

/// Recomputes the neighbours once ready, if the member list, the reachable members or
/// `k` changed.
pub(crate) async fn refresh(state: &AppState) -> Option<NeighbourChange> {
    if !state.is_ready().await {
        return None;
    }
    let (current, stable) = {
        let replica = state.doc.read().await;
        let current = members::members(&replica.doc);
        let stable: BTreeSet<String> = current
            .iter()
            .filter(|m| members::tier(&replica.doc, m) == members::Tier::Stable)
            .cloned()
            .collect();
        (current, stable)
    };
    let k = state.neighbours_k;
    let now = Instant::now();
    let mut unreachable = BTreeSet::new();
    for member in &current {
        if let Some(peer_state) = state.get_peer_state(member).await {
            if peer_state.lock().await.is_unreachable(&state.timing, now) {
                unreachable.insert(member.clone());
            }
        }
    }
    let previous = {
        let view = state
            .neighbours
            .read()
            .unwrap_or_else(|poisoned| poisoned.into_inner());
        if view.as_ref().is_some_and(|view| {
            view.members == current
                && view.k == k
                && view.unreachable == unreachable
                && view.stable == stable
        }) {
            return None;
        }
        view.as_ref().map(|view| view.neighbours.clone())
    };

    // Before the first view everyone known counted as a neighbour.
    let previous = match previous {
        Some(previous) => previous,
        None => state.peer_states.read().await.keys().cloned().collect(),
    };
    let computed = neighbours(
        &current,
        state.local_id(),
        k,
        &|m| unreachable.contains(m),
        &|m| stable.contains(m),
    );
    let change = NeighbourChange {
        added: computed
            .iter()
            .filter(|m| !previous.contains(*m))
            .cloned()
            .collect(),
        removed: previous
            .iter()
            .filter(|m| !computed.contains(*m))
            .cloned()
            .collect(),
    };
    info!(
        members = current.len(),
        unreachable = unreachable.len(),
        stable = stable.len(),
        k,
        neighbours = computed.len(),
        added = change.added.len(),
        removed = change.removed.len(),
        "Neighbours computed from the member list"
    );
    let reachable_others = current
        .iter()
        .filter(|m| m.as_str() != state.local_id() && !unreachable.contains(*m))
        .count();
    *state
        .neighbours
        .write()
        .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(NeighbourView {
        full_mesh: computed.len() >= reachable_others,
        members: current,
        k,
        unreachable,
        stable,
        neighbours: computed,
    });
    Some(change)
}

const RING_LABEL: &[u8] = b"continuum/topology-ring/1";
const PAIR_LABEL: &[u8] = b"continuum/topology-pair/1";

/// 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. Costs about N² hashes (a few milliseconds for
/// hundreds of members). `stable` members (`swarm::stability`) are preferred for half of
/// each member's scored slots.
///
/// 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.
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")
    };

    for member in &live {
        let at = position[member];
        let next = first_live(at, 1);
        let prev = first_live(at, n - 1);
        link(member, next);
        link(member, prev);
        let mut candidates: Vec<([u8; 32], &str)> = live
            .iter()
            .copied()
            .filter(|other| other != member && *other != next && *other != prev)
            .map(|other| (pair_score(member, other), other))
            .collect();
        candidates.sort_unstable();
        // 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 and the load
        // on stable members bounded.
        let stable_slots = (k - 2).div_ceil(2);
        let preferred: Vec<&str> = candidates
            .iter()
            .map(|(_, other)| *other)
            .filter(|other| stable(other))
            .take(stable_slots)
            .collect();
        let rest = candidates
            .iter()
            .map(|(_, other)| *other)
            .filter(|other| !preferred.contains(other))
            .take(k - 2 - preferred.len());
        for partner in preferred.iter().copied().chain(rest) {
            link(member, partner);
        }
    }
    links
}

/// Symmetric: the same for (a, b) and (b, a).
fn pair_score(a: &str, b: &str) -> [u8; 32] {
    let (low, high) = if a <= b { (a, b) } else { (b, a) };
    hash(PAIR_LABEL, &[low, high])
}

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()
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::VecDeque;

    fn members(n: usize) -> Vec<String> {
        (0..n).map(|i| format!("{:064x}", i * 7919 + 13)).collect()
    }

    fn all_reachable(_: &str) -> bool {
        false
    }

    fn none_stable(_: &str) -> bool {
        false
    }

    fn graph(members: &[String], k: usize) -> BTreeMap<String, BTreeSet<String>> {
        super::graph(members, k, &all_reachable, &none_stable)
    }

    fn connected(g: &BTreeMap<String, BTreeSet<String>>) -> bool {
        let Some(start) = g.keys().next() else {
            return true;
        };
        let mut seen = BTreeSet::from([start.clone()]);
        let mut queue = VecDeque::from([start.clone()]);
        while let Some(m) = queue.pop_front() {
            for next in &g[&m] {
                if seen.insert(next.clone()) {
                    queue.push_back(next.clone());
                }
            }
        }
        seen.len() == g.len()
    }

    fn edges(graph: &BTreeMap<String, BTreeSet<String>>) -> BTreeSet<(String, String)> {
        graph
            .iter()
            .flat_map(|(a, ns)| {
                ns.iter().map(move |b| {
                    if a < b {
                        (a.clone(), b.clone())
                    } else {
                        (b.clone(), a.clone())
                    }
                })
            })
            .collect()
    }

    #[test]
    fn small_swarms_are_a_full_mesh() {
        let ms = members(13);
        for m in &ms {
            assert_eq!(
                neighbours(&ms, m, 12, &all_reachable, &none_stable).len(),
                12
            );
        }
    }

    #[test]
    fn links_are_symmetric_and_degree_at_least_k() {
        let ms = members(250);
        let g = graph(&ms, 8);
        for (a, ns) in &g {
            assert!(ns.len() >= 8, "{a} has {} neighbours", ns.len());
            assert!(!ns.contains(a));
            for b in ns {
                assert!(g[b].contains(a), "{a} -> {b} not symmetric");
            }
        }
    }

    #[test]
    fn graph_is_connected() {
        for (n, k) in [(250, 8), (250, 2), (60, 3)] {
            assert!(connected(&graph(&members(n), k)), "n={n} k={k}");
        }
    }

    #[test]
    fn one_member_joining_changes_few_links() {
        let ms = members(200);
        let before = edges(&graph(&ms, 8));
        let mut more = ms.clone();
        more.push(format!("{:064x}", 999_999));
        let after = edges(&graph(&more, 8));
        let changed = before.symmetric_difference(&after).count();
        assert!(changed <= 4 * 8, "{changed} links changed");
    }

    #[test]
    fn reachable_members_stay_connected_when_most_are_offline() {
        let ms = members(100);
        // Every tenth member is online.
        let online: BTreeSet<String> = ms.iter().step_by(10).cloned().collect();
        let offline = |m: &str| !online.contains(m);
        let g = super::graph(&ms, 12, &offline, &none_stable);
        assert_eq!(g.len(), 10);
        // Ten reachable members and k = 12: a full mesh among them.
        assert!(g.values().all(|ns| ns.len() == 9));

        let online: BTreeSet<String> = ms.iter().step_by(4).cloned().collect();
        let offline = |m: &str| !online.contains(m);
        let g = super::graph(&ms, 8, &offline, &none_stable);
        assert_eq!(g.len(), 25);
        assert!(connected(&g));
        assert!(
            g.values()
                .all(|ns| ns.len() >= 8 && ns.iter().all(|n| online.contains(n)))
        );
    }

    #[test]
    fn stable_members_become_hubs_with_bounded_degree() {
        let ms = members(250);
        let stable: BTreeSet<String> = ms.iter().step_by(12).take(20).cloned().collect();
        let is_stable = |m: &str| stable.contains(m);
        let g = super::graph(&ms, 12, &all_reachable, &is_stable);
        assert!(connected(&g));

        let degree = |pred: &dyn Fn(&String) -> bool| -> (usize, f64) {
            let degrees: Vec<usize> = g
                .iter()
                .filter(|(m, _)| pred(m))
                .map(|(_, ns)| ns.len())
                .collect();
            let max = *degrees.iter().max().unwrap();
            (
                max,
                degrees.iter().sum::<usize>() as f64 / degrees.len() as f64,
            )
        };
        let (stable_max, stable_mean) = degree(&|m| stable.contains(m));
        let (_, transient_mean) = degree(&|m| !stable.contains(m));
        // 230 transient members put about 5 of their 10 scored slots on 20 stable ones.
        assert!(
            stable_mean > 3.0 * transient_mean,
            "{stable_mean} vs {transient_mean}"
        );
        assert!(
            stable_max <= 100,
            "a stable member has {stable_max} neighbours"
        );
    }
}