//! 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 same graph from the member list alone, with no coordination:
//! * a ring ordered by key hash: each member's predecessor and successor, which keeps
//!   the graph connected;
//! * plus, for each member, the `k - 2` 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` members the graph is the full mesh: small swarms
//! behave exactly as before. A member joining or leaving changes only links near it.

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

use sha2::{Digest, Sha256};
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,
    neighbours: BTreeSet<String>,
}

/// 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 (or `k`) changed.
pub(crate) async fn refresh(state: &AppState) -> Option<NeighbourChange> {
    if !state.is_ready().await {
        return None;
    }
    let current = members::members(&state.doc.read().await.doc);
    let k = state.neighbours_k;
    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)
        {
            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);
    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(),
        k,
        neighbours = computed.len(),
        added = change.added.len(),
        removed = change.removed.len(),
        "Neighbours computed from the member list"
    );
    *state
        .neighbours
        .write()
        .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(NeighbourView {
        members: current,
        k,
        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`).
pub(crate) fn neighbours(members: &[String], local: &str, k: usize) -> BTreeSet<String> {
    let mut all: Vec<String> = members.to_vec();
    all.push(local.to_owned());
    graph(&all, k).remove(local).unwrap_or_default()
}

/// Every member's neighbours. Costs about N² hashes (a few milliseconds for hundreds of
/// members); computed again only when the member list changes.
pub(crate) fn graph(members: &[String], k: usize) -> BTreeMap<String, BTreeSet<String>> {
    let k = k.max(2);
    let all: BTreeSet<&str> = members.iter().map(String::as_str).collect();
    let mut links: BTreeMap<String, BTreeSet<String>> = all
        .iter()
        .map(|m| ((*m).to_owned(), BTreeSet::new()))
        .collect();

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

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

    let ring = ring_links(&all);
    for (member, [prev, next]) in &ring {
        link(member, prev);
        link(member, next);
        // Its k - 2 partners with the lowest pair score, besides its ring links.
        let mut candidates: Vec<([u8; 32], &str)> = all
            .iter()
            .copied()
            .filter(|other| other != member && other != prev && other != next)
            .map(|other| (pair_score(member, other), other))
            .collect();
        candidates.sort_unstable();
        for (_, partner) in candidates.into_iter().take(k - 2) {
            link(member, partner);
        }
    }
    links
}

/// Ring predecessor and successor of every member, by key hash.
fn ring_links<'a>(all: &BTreeSet<&'a str>) -> BTreeMap<&'a str, [&'a str; 2]> {
    let mut ring: Vec<([u8; 32], &'a str)> =
        all.iter().map(|m| (hash(RING_LABEL, &[*m]), *m)).collect();
    ring.sort_unstable();
    let n = ring.len();
    (0..n)
        .map(|i| (ring[i].1, [ring[(i + n - 1) % n].1, ring[(i + 1) % n].1]))
        .collect()
}

/// 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 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).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)] {
            let ms = members(n);
            let g = graph(&ms, k);
            let mut seen = BTreeSet::from([ms[0].clone()]);
            let mut queue = VecDeque::from([ms[0].clone()]);
            while let Some(m) = queue.pop_front() {
                for next in &g[&m] {
                    if seen.insert(next.clone()) {
                        queue.push_back(next.clone());
                    }
                }
            }
            assert_eq!(seen.len(), n, "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");
    }
}