//! PlumTree (Leitão, Pereira and Rodrigues, "Epidemic Broadcast Trees", 2007) as the way
//! changes reach every member (`swarm::dissemination`, `--dissemination plumtree`), over
//! the same neighbours as writers' trees (`swarm::topology`).
//!
//! Each node keeps which of its neighbours it pushes changes to at once (eager) and
//! which it leaves to announcements (lazy), one set for all writers, as in the paper:
//! at first all are eager. A change that comes along a writer's push with nothing new
//! (a duplicate: it came another way first) makes the receiver ask the sender to leave
//! changes to announcements from now on (PRUNE), and both hold the link lazy. A change
//! that comes another way (pulled after an announcement, PlumTree's IHAVE: the eager
//! links didn't bring it) makes the receiver ask the one it came from to push at once
//! (GRAFT), and both hold the link eager. So the eager links settle into a tree made of
//! the links that delivered first: one that follows latency by itself, where writers'
//! trees are computed alike on every node from the member list (`swarm::tree`). One
//! tree for all writers settles over every writer's changes; trees per writer (tried
//! first) never did with writers in turn, each writing too rarely, every new one
//! flooding until pruned. PRUNE and GRAFT name the writer, for such a variant; this one
//! doesn't use it.
//!
//! PlumTree's IHAVE is the announcement (`swarm::relay`): lazy neighbours, a few of them
//! at random (`LAZY_FANOUT`), are told the new heads in the quick round, coalesced, all
//! of them in the slow round, and count as having the changes
//! (`Plan::covered`), as their eager links bring them. One that still lacks them a round
//! trip later pulls (`relay::pull_after_grace`: the graft timer), and the pull, bringing
//! a writer's changes another way, grafts the link. Told only in the slow round instead
//! (tried second), repair took seconds; every lazy link told per round is PlumTree's
//! cost (one STATUS each; a message of its own would only be smaller), so telling all
//! in the quick round tripled the messages (tried third).
//!
//! The paper's optimisation: a lazy neighbour that told us of changes before an eager
//! one pushed them has the faster link (an announcement waits `relay::ANNOUNCE_DELAY`,
//! a push doesn't), so it is grafted, once it was so for a few pushes in a row
//! (`AHEAD_TIMES`); the old parent's pushes then come as duplicates,
//! and are pruned. A link we just grafted isn't pruned for a duplicate while what it
//! pushed meanwhile may still be on its way (`settling`). Else a tree stays as it first
//! formed: a pull beats a push only when
//! the push is slower by an announcement's delay and two round trips.
//! Nothing is written to the document.

use std::collections::{HashMap, HashSet};
use std::sync::Mutex;
use std::time::Duration;

use async_trait::async_trait;
use tokio::time::Instant;
use tracing::debug;

use crate::document::ChangeHash;
use crate::locks::MutexExt;
use crate::state::AppState;
use crate::swarm::dissemination::{Dissemination, pass_on};
use crate::swarm::topology;
use crate::swarm::tree::Plan;
use crate::transport::{ParsedMessage, TreeLink};

/// Lazy links kept at most; past it, all are forgotten (every link eager again, and
/// pruned again where it delivers twice). More than a node has neighbours.
const MAX_LAZY: usize = 4096;

/// Lazy neighbours told in the quick round per change passed on, picked at random each
/// time (`some_of`); the others hear in the slow round. Told all (patch 997), every node
/// announced to all its lazy neighbours at every hop: about 240 STATUS per write at 50
/// nodes, three times the messages of PlumTree told only in the slow round.
const LAZY_FANOUT: usize = 3;

/// Up to `count` of `peers`, picked at random.
fn some_of(peers: &[String], count: usize) -> Vec<String> {
    let mut peers = peers.to_vec();
    let count = count.min(peers.len());
    for picked in 0..count {
        let rest = (peers.len() - picked) as u64;
        let other = picked + (crate::transport::paths::random_nonce() % rest) as usize;
        peers.swap(picked, other);
    }
    peers.truncate(count);
    peers
}

/// Heads we lacked when told of them, kept at most; past it, all are forgotten (a
/// better link is noticed with the next change).
const MAX_ANNOUNCED: usize = 4096;

/// Who first told us of a head we lacked is noticed for this long.
const ANNOUNCED_FOR: Duration = Duration::from_secs(10);

/// Pushes in a row a lazy neighbour must have told of first to be grafted (`ahead`).
/// Grafted at once (patch 1001), trees shifted with every lucky announcement: on Tor,
/// with one writer, 24 PRUNE and GRAFT per write against 6, and changes along the
/// eager links took 987 ms against 471.
const AHEAD_TIMES: u32 = 3;

/// How long after we grafted a link its duplicates don't prune it, in round trips to
/// the peer (`settling`), and when we know none.
const SETTLE_ROUND_TRIPS: u32 = 2;
const SETTLE_UNKNOWN: Duration = Duration::from_secs(2);

/// The neighbours whose links are lazy, who first told us of heads we lacked, how many
/// pushes in a row a lazy one told of first, and when we grafted each.
#[derive(Default)]
pub(crate) struct PlumTree {
    lazy: Mutex<HashSet<String>>,
    announced: Mutex<HashMap<ChangeHash, (String, Instant)>>,
    ahead_times: Mutex<HashMap<String, u32>>,
    /// When we last grafted each link (`settling`).
    grafted: Mutex<HashMap<String, Instant>>,
}

impl PlumTree {
    /// Holds the link to `peer` lazy, or eager; whether that changed it.
    fn set(&self, peer: &str, lazy: bool) -> bool {
        let mut links = self.lazy.lock_or_recover();
        if lazy {
            if links.len() >= MAX_LAZY {
                links.clear();
            }
            links.insert(peer.to_owned())
        } else {
            links.remove(peer)
        }
    }

    /// A lazy neighbour, not `peer`, that told us of any of the changes `peer` just
    /// pushed before they came (since `ANNOUNCED_FOR`), for the last `AHEAD_TIMES`
    /// pushes in a row: its link is the faster, by at least an announcement's delay
    /// (`relay::ANNOUNCE_DELAY`), which a push over it doesn't wait. The changes are
    /// forgotten as announced.
    fn ahead(&self, peer: &str, authors: &[(String, Vec<ChangeHash>)]) -> Option<String> {
        let now = Instant::now();
        let told: Vec<(String, Instant)> = {
            let mut announced = self.announced.lock_or_recover();
            authors
                .iter()
                .flat_map(|(_, changes)| changes)
                .filter_map(|change| announced.remove(change))
                .collect()
        };
        let first = {
            let lazy = self.lazy.lock_or_recover();
            told.into_iter()
                .filter(|(by, at)| {
                    by != peer && lazy.contains(by) && now.duration_since(*at) < ANNOUNCED_FOR
                })
                .min_by_key(|(_, at)| *at)
                .map(|(by, _)| by)
        };
        let mut times = self.ahead_times.lock_or_recover();
        let Some(first) = first else {
            times.clear();
            return None;
        };
        times.retain(|by, _| *by == first);
        let count = times.entry(first.clone()).or_default();
        *count += 1;
        if *count < AHEAD_TIMES {
            return None;
        }
        times.clear();
        Some(first)
    }

    /// We grafted the link to `peer` now.
    fn grafted(&self, peer: &str) {
        let mut grafted = self.grafted.lock_or_recover();
        if grafted.len() >= MAX_LAZY {
            grafted.clear();
        }
        grafted.insert(peer.to_owned(), Instant::now());
    }

    /// Whether we grafted the link to `peer` so lately (`SETTLE_ROUND_TRIPS` round trips
    /// to it) that a duplicate over it doesn't prune it: what it pushed was on its way
    /// as we pulled the same from it, or from another. Pruned at once, a node no
    /// neighbour pushed to pulled every change, grafted the one it pulled from, and
    /// pruned it at its next push, a duplicate of the pull, over and over (patch 1007:
    /// on Tor, such a node got every change late).
    fn settling(&self, state: &AppState, peer: &str) -> bool {
        let Some(at) = self.grafted.lock_or_recover().get(peer).copied() else {
            return false;
        };
        let now = Instant::now();
        let settle = state
            .lock_rtt()
            .effective(peer, now, state.timing.rtt_fresh_for())
            .map_or(SETTLE_UNKNOWN, |rtt| rtt * SETTLE_ROUND_TRIPS);
        now.duration_since(at) < settle
    }

    /// The changes, come another way, are forgotten as announced.
    fn forget_announced(&self, authors: &[(String, Vec<ChangeHash>)]) {
        let mut announced = self.announced.lock_or_recover();
        for change in authors.iter().flat_map(|(_, changes)| changes) {
            announced.remove(change);
        }
    }

    /// Tells `peer` to hold the link for `writer` as we now do: a GRAFT (`eager`) or a
    /// PRUNE.
    async fn tell(state: &AppState, peer: &str, writer: u64, eager: bool) {
        let (epoch, epoch_seq) = state.epoch_stamp();
        let link = TreeLink {
            epoch,
            epoch_seq,
            writer,
        };
        let message = if eager {
            ParsedMessage::Graft(link)
        } else {
            ParsedMessage::Prune(link)
        };
        if let Err(error) = state.transport.send_message(peer, &message).await {
            debug!(peer = %peer, error = %format!("{error:#}"), "PRUNE or GRAFT failed");
        }
    }
}

#[async_trait]
impl Dissemination for PlumTree {
    fn name(&self) -> &'static str {
        "plumtree"
    }

    /// Eager neighbours get the writer's changes at once (`Plan::children`); lazy ones
    /// count as having them (`Plan::covered`), `LAZY_FANOUT` of them, picked at random,
    /// told in the quick round (`Plan::uncovered`: the IHAVE), all in the slow round.
    async fn plan(&self, state: &AppState, writer: u64) -> Option<Plan> {
        let _ = writer;
        let neighbours = topology::our_neighbours(state);
        let lazy = self.lazy.lock_or_recover().clone();
        let (lazy, eager): (Vec<String>, Vec<String>) =
            neighbours.into_iter().partition(|peer| lazy.contains(peer));
        Some(Plan {
            children: eager,
            alone: Vec::new(),
            uncovered: some_of(&lazy, LAZY_FANOUT),
            covered: lazy,
        })
    }

    async fn forward(
        &self,
        state: &AppState,
        writer: u64,
        from: &str,
        changes: &[ChangeHash],
    ) -> Option<Plan> {
        let plan = self.plan(state, writer).await?;
        Some(pass_on(state, writer, from, changes, plan).await)
    }

    /// A duplicate along a writer's push prunes the link; changes of a writer that
    /// came another way graft the link to whoever brought them. Told every time, not
    /// only when our side of the link changed: a duplicate pushed over a link we hold
    /// lazy, or changes pulled over one we hold eager, say the other end holds it the
    /// other way (a PRUNE or GRAFT lost, or two crossing), and told only on a change,
    /// it stayed so: a node every neighbour held lazy while it held them eager pulled
    /// every change (patch 999: one such on Tor, a few links eager at one end). Counted as the trees'
    /// counters count (`/admin/stats`): a duplicate as a message from another than the
    /// parent, changes another way as ones the trees didn't bring.
    async fn note_arrival(
        &self,
        state: &AppState,
        peer: &str,
        writer: Option<u64>,
        carried: bool,
        authors: &[(String, Vec<ChangeHash>)],
    ) {
        match writer {
            Some(writer) if carried && authors.is_empty() => {
                state.stats.tree_from_other();
                if self.settling(state, peer) {
                    return;
                }
                self.set(peer, true);
                Self::tell(state, peer, writer, false).await;
            }
            Some(writer) if !authors.is_empty() => {
                self.set(peer, false);
                if let Some(ahead) = self.ahead(peer, authors) {
                    debug!(peer = %ahead, "A lazy neighbour told of changes before they were pushed, every time lately; grafting it");
                    self.set(&ahead, false);
                    self.grafted(&ahead);
                    Self::tell(state, &ahead, writer, true).await;
                }
            }
            None => {
                self.forget_announced(authors);
                for (author, changes) in authors {
                    if author == state.local_id() {
                        continue;
                    }
                    state.stats.tree_missed(changes.len());
                    self.set(peer, false);
                    self.grafted(peer);
                    Self::tell(state, peer, topology::pair_key(author), true).await;
                }
            }
            Some(_) => {}
        }
    }

    async fn linked(&self, state: &AppState, peer: &str, writer: u64, eager: bool) {
        let _ = (state, writer);
        self.set(peer, !eager);
    }

    /// The first to tell us of each head is the one noticed.
    fn announced(&self, peer: &str, heads: &[ChangeHash]) {
        let mut announced = self.announced.lock_or_recover();
        if announced.len() + heads.len() > MAX_ANNOUNCED {
            announced.clear();
        }
        let now = Instant::now();
        for head in heads {
            announced
                .entry(*head)
                .or_insert_with(|| (peer.to_owned(), now));
        }
    }

    #[cfg(test)]
    fn lazy(&self) -> Option<HashSet<String>> {
        Some(self.lazy.lock_or_recover().clone())
    }
}