//! Passing a writer's changes on along its tree: every member gets them once, over a
//! link that works, and as soon as the member before it has them.
//!
//! The tree is the same on every node, computed from what they share: the graph the
//! member list gives (`topology::member_list_graph`: every member's neighbours, as if
//! all reached all) and what each member tells its neighbours it reaches
//! (`swarm::reach`). Every member's distance from the writer over that graph orders the
//! members: by distance, then whether it reaches a member a hop closer, then pair key;
//! a member's parent is the neighbour it reaches that comes before it, the closest to
//! the writer first, then one that reaches closer, then by pair score
//! (`topology::pair_score`). Coming before rules out cycles, and taking only a neighbour
//! it reaches routes around links that don't work: where a writer behind NAT reaches a
//! single member, that one comes first and the others take it.
//!
//! Whose a change is, its author says (every node writes as itself: its node ID, by
//! pair key here): whoever applies changes passes them on to its own children in each of
//! their writers' trees (`forward`), however they came; so does the writer, to all its
//! neighbours, with its own writes. A sync message along a tree also says whose
//! (`SyncEnvelope::writer`), for one that had them already (`swarm::sync`), and needs no
//! answer where the child had nothing the parent lacks (`SyncSession::assume_delivered`).
//! A member no member it reaches comes before (or whose parent didn't get the change)
//! gets it from the announcements (`swarm::relay`), as before.
//!
//! The graph is computed once per member list (tens of milliseconds for thousands of
//! members), the distances once per writer and graph (`MAX_WRITERS` kept).

use std::collections::{BTreeMap, BTreeSet, HashMap, VecDeque};
use std::sync::Arc;

use tracing::trace;

use crate::locks::MutexExt;
use crate::state::AppState;
use crate::swarm::{reach, topology};

/// The writers whose distances are kept at most.
const MAX_WRITERS: usize = 256;

type Graph = BTreeMap<String, BTreeSet<String>>;

/// The distances from each writer, over the graph they were computed for.
#[derive(Default)]
pub(crate) struct Trees {
    graph: Option<Arc<Graph>>,
    /// The members by pair key, of that graph.
    by_key: HashMap<u64, String>,
    distances: HashMap<u64, Arc<HashMap<String, u32>>>,
}

/// The graph, and every member's distance from `writer` (by pair key) over it; None
/// while the member list isn't known, or `writer` isn't in it.
fn distances(state: &AppState, writer: u64) -> Option<(Arc<Graph>, Arc<HashMap<String, u32>>)> {
    let graph = topology::member_list_graph(state)?;
    let mut trees = state.trees.lock_or_recover();
    if !trees
        .graph
        .as_ref()
        .is_some_and(|known| Arc::ptr_eq(known, &graph))
    {
        trees.by_key = graph
            .keys()
            .map(|member| (topology::pair_key(member), member.clone()))
            .collect();
        trees.distances.clear();
        trees.graph = Some(Arc::clone(&graph));
    }
    if let Some(known) = trees.distances.get(&writer) {
        return Some((graph, Arc::clone(known)));
    }
    let from = trees.by_key.get(&writer)?.clone();
    let mut distance = HashMap::from([(from.clone(), 0u32)]);
    let mut queue = VecDeque::from([from]);
    while let Some(member) = queue.pop_front() {
        let next = distance[&member] + 1;
        for neighbour in graph.get(&member).into_iter().flatten() {
            if !distance.contains_key(neighbour) {
                distance.insert(neighbour.clone(), next);
                queue.push_back(neighbour.clone());
            }
        }
    }
    if trees.distances.len() >= MAX_WRITERS {
        trees.distances.clear();
    }
    let distance = Arc::new(distance);
    trees.distances.insert(writer, Arc::clone(&distance));
    Some((graph, distance))
}

/// Our children in `writer`'s tree: our neighbours in the graph whose parent we are,
/// and whom we reach.
///
/// Members are ordered by distance from the writer, then whether they reach a member a
/// hop closer (as they say: else, while it hasn't the change from its own side, it is
/// as good as farther), then pair key; a member's parent is the neighbour it reaches
/// that comes before it, by the same, then pair score. So where the writer reaches few
/// (behind NAT), those it reaches come first among the members as far from it, and the
/// others take them as parents. What a member reaches we know of ourselves and our
/// neighbours, as they said; of others, the graph is taken as right (if we're wrong,
/// a change goes twice or by announcement).
#[cfg(test)]
pub(crate) async fn children(state: &AppState, writer: u64) -> Vec<String> {
    plan(state, writer)
        .await
        .map(|plan| plan.children)
        .unwrap_or_default()
}

/// Where `member` comes in the writer's tree (by `distance` from it over `graph`):
/// its distance, then whether it reaches a member a hop closer (`reached(by, other)`:
/// whether `by` reaches `other`, as known; the writer does), then its pair key. None if
/// the graph doesn't have it.
fn rank(
    graph: &Graph,
    distance: &HashMap<String, u32>,
    reached: &dyn Fn(&str, &str) -> bool,
    member: &str,
) -> Option<(u32, bool, u64)> {
    let d = *distance.get(member)?;
    let closer =
        d == 0
            || graph.get(member).into_iter().flatten().any(|other| {
                distance.get(other.as_str()) == Some(&(d - 1)) && reached(member, other)
            });
    Some((d, !closer, topology::pair_key(member)))
}

/// `member`'s parent in the writer's tree (`rank`): the neighbour it reaches that comes
/// before it, the closest to the writer first, then one that reaches closer, then by
/// pair score; None if none does (or it is the writer).
fn parent<'g>(
    graph: &'g Graph,
    distance: &HashMap<String, u32>,
    reached: &dyn Fn(&str, &str) -> bool,
    member: &str,
) -> Option<&'g str> {
    let its_rank = rank(graph, distance, reached, member)?;
    graph
        .get(member)
        .into_iter()
        .flatten()
        .map(String::as_str)
        .filter(|candidate| {
            rank(graph, distance, reached, candidate).is_some_and(|r| r < its_rank)
                && reached(member, candidate)
        })
        .min_by_key(|candidate| {
            let (d, far, _) = rank(graph, distance, reached, candidate).expect("ranked");
            (d, far, topology::pair_score(candidate, member))
        })
}

/// `writer`'s whole tree, as far as `reached` knows whom each member reaches (the
/// dashboard's map, from the links members report): every member's parent, None for
/// those without one (they get the writer's changes by announcement). None while the
/// member list isn't known, or the writer isn't in it.
#[cfg_attr(not(feature = "dashboard"), allow(dead_code))]
pub(crate) fn whole(
    state: &AppState,
    writer: &str,
    reached: &dyn Fn(&str, &str) -> bool,
) -> Option<BTreeMap<String, Option<String>>> {
    let (graph, distance) = distances(state, topology::pair_key(writer))?;
    Some(
        graph
            .keys()
            .filter(|member| distance.get(member.as_str()).is_some_and(|d| *d > 0))
            .map(|member| {
                let parent = parent(&graph, &distance, reached, member).map(str::to_owned);
                (member.clone(), parent)
            })
            .collect(),
    )
}

/// What `writer`'s tree means for us: our children in it, and our neighbours it leaves
/// out (`uncovered`).
pub(crate) struct Plan {
    pub(crate) children: Vec<String>,
    /// Those of our children that reach no member but us, as they said: nobody else's
    /// slow round tells them of a change, should our message to them be lost (it is
    /// taken as delivered, `SyncSession::assume_delivered`), so ours does.
    pub(crate) alone: Vec<String>,
    /// Our neighbours that no member they reach comes before (the tree reaches them only
    /// by announcement), and those not in the member list's graph (whom no tree has).
    pub(crate) uncovered: Vec<String>,
}

/// Our part in `writer`'s tree (`children`, `Plan`); None while the member list isn't
/// known, or the writer or we aren't in it.
pub(crate) async fn plan(state: &AppState, writer: u64) -> Option<Plan> {
    let (graph, distance) = distances(state, writer)?;
    let me = state.local_id();
    if !distance.contains_key(me) {
        return None;
    }
    // Whom we and our neighbours reach, by pair key.
    let mut reaches: HashMap<&str, Vec<u64>> = HashMap::new();
    reaches.insert(me, reach::published(state));
    for member in graph.get(me).into_iter().flatten() {
        if let Some(peer_state) = state.get_peer_state(member).await
            && let Some(keys) = reach::of(peer_state.lock().await.reaches.as_ref())
        {
            reaches.insert(member.as_str(), keys);
        }
    }
    let reached = |by: &str, member: &str| {
        reaches
            .get(by)
            .is_none_or(|keys| keys.contains(&topology::pair_key(member)))
    };
    let my_rank = rank(&graph, &distance, &reached, me)?;
    let my_key = topology::pair_key(me);
    let mut children = Vec::new();
    let mut alone = Vec::new();
    let mut uncovered = Vec::new();
    for member in graph.get(me).into_iter().flatten() {
        let Some(its_rank) = rank(&graph, &distance, &reached, member) else {
            continue;
        };
        if distance[member.as_str()] == 0 {
            continue;
        }
        match parent(&graph, &distance, &reached, member) {
            None => uncovered.push(member.clone()),
            Some(parent)
                if parent == me && its_rank > my_rank && topology::reaches(state, member).await =>
            {
                if reaches
                    .get(member.as_str())
                    .is_some_and(|keys| keys.iter().all(|key| *key == my_key))
                {
                    alone.push(member.clone());
                }
                children.push(member.clone());
            }
            Some(_) => {}
        }
    }
    // Neighbours the member list's graph doesn't give us (not listed yet, or only
    // connecting out): no tree has them.
    for (peer, _) in state.peers_snapshot().await {
        if topology::is_neighbour(state, &peer)
            && !graph.get(me).is_some_and(|ours| ours.contains(&peer))
        {
            uncovered.push(peer);
        }
    }
    Some(Plan {
        children,
        alone,
        uncovered,
    })
}

/// We applied a change of `writer`'s that `from` passed on: to our children in its
/// tree (but `from`), a sync message each, saying whose. Returns our part in the tree
/// (`Plan`), None if we have none (announcements then do it all).
pub(crate) async fn forward(state: &AppState, writer: u64, from: &str) -> Option<Plan> {
    let mut plan = plan(state, writer).await?;
    plan.children.retain(|child| child != from);
    let children = &plan.children;
    for child in children {
        if let Some(peer_state) = state.get_peer_state(child).await {
            peer_state.lock().await.tree_writer = Some(writer);
        }
        state.stats.tree_forward();
        state.request_sync(child);
    }
    if !children.is_empty() {
        trace!(
            writer,
            children = children.len(),
            "Passing a change on along the writer's tree"
        );
    }
    Some(plan)
}