//! Keepalives (PING / PONG) and the self-healing re-handshake when they show our
//! documents differ while sync is idle.

use std::sync::atomic::Ordering;
use std::time::Duration;

use anyhow::Result;
use tokio::time::Instant;
use tracing::warn;

use super::handshake::ensure_handshake;
use crate::state::{AppState, Liveness, Standing};
use crate::swarm::sync::SyncEvent;
use crate::swarm::topology;
use crate::transport::{self, ParsedMessage};

/// Sends `peer` a PING with our keepalive, noting the time for its round trip.
pub(crate) async fn send_ping(state: &AppState, peer: &str) -> Result<()> {
    let keepalive = local_keepalive(state, peer).await;
    state.lock_rtt().ping_sent(peer, Instant::now());
    state
        .transport
        .send_message(peer, &ParsedMessage::Ping(keepalive))
        .await
}

/// How long shutting down waits for the last PINGs (`farewell`).
const FAREWELL_TIMEOUT: Duration = Duration::from_secs(2);

/// At shutdown: one last PING to each neighbour, saying we aren't ready. A neighbour
/// catching up from us picks another source now, instead of once we go dormant. Bounded
/// in time: our transport may be gone already.
pub(crate) async fn farewell(state: &AppState) {
    state.stopping.store(true, Ordering::Release);
    let mut pings = tokio::task::JoinSet::new();
    for (peer, _) in state.peers_snapshot().await {
        if topology::is_neighbour(state, &peer) {
            let state = state.clone();
            pings.spawn(async move { send_ping(&state, &peer).await });
        }
    }
    let all_sent = async { while pings.join_next().await.is_some() {} };
    let _ = tokio::time::timeout(FAREWELL_TIMEOUT, all_sent).await;
}

/// Our keepalive for `peer`.
pub(crate) async fn local_keepalive(state: &AppState, peer: &str) -> transport::Keepalive {
    // Heads and epoch are published together (`Replica::stamps`), so they always match,
    // and reading them doesn't wait for a sync message being applied.
    let (heads, epoch, epoch_seq) = {
        let stamp = state.stamp.borrow();
        (
            transport::HeadsSummary::of(stamp.heads.to_vec()),
            stamp.epoch,
            stamp.epoch_seq,
        )
    };
    let syncs_with_you = match state.get_peer_state(peer).await {
        Some(peer_state) => {
            let peer_state = peer_state.lock().await;
            peer_state.liveness == Liveness::Active
                && matches!(peer_state.standing, Standing::Member { .. })
                && peer_state.epoch == Some(epoch)
        }
        None => false,
    };
    transport::Keepalive {
        heads,
        epoch,
        epoch_seq,
        ready: state.is_ready().await && !state.stopping.load(Ordering::Acquire),
        syncs_with_you,
    }
}

/// What a peer's keepalive says about it.
pub(super) async fn note_keepalive(state: &AppState, peer: &str, remote: &transport::Keepalive) {
    if let Some(peer_state) = state.get_peer_state(peer).await {
        let mut peer_state = peer_state.lock().await;
        let became_ready = remote.ready && peer_state.reports_ready != Some(true);
        peer_state.reports_ready = Some(remote.ready);
        peer_state.syncs_with_us = Some(remote.syncs_with_you);
        peer_state.reported_heads = remote.heads.heads().map(<[_]>::len);
        drop(peer_state);
        // Off the consumer: telling which heads we lack takes the document's lock.
        {
            let (state, peer, remote) = (state.clone(), peer.to_owned(), remote.clone());
            tokio::spawn(async move {
                crate::swarm::relay::pull_if_behind(&state, &peer, &remote).await;
            });
        }
        if became_ready {
            // We may have left the opening to it while it wasn't ready, and nothing else
            // would make us look again.
            let _ = state
                .sync_events
                .send(SyncEvent::PeerNeedsSync(peer.to_owned()));
        }
    }
    crate::swarm::source::reported(state, peer, remote.ready, Some(remote.syncs_with_you));
}

pub(super) async fn heal_if_diverged(
    state: &AppState,
    peer: &str,
    local: &transport::Keepalive,
    remote: &transport::Keepalive,
) {
    // Heads of documents from different epochs are never equal and never converge.
    // And a peer that isn't our neighbour isn't ours to keep in step.
    if local.epoch != remote.epoch || !topology::is_neighbour(state, peer) {
        return;
    }

    let Some(peer_state_mutex) = state.get_peer_state(peer).await else {
        return;
    };

    let heal = {
        let mut peer_state = peer_state_mutex.lock().await;
        if local.heads == remote.heads {
            // Same document: whatever answer we were waiting for, there is nothing left
            // to send. (Automerge answers only when it has something to say, so a sync
            // message to a peer that already has everything gets no reply at all.)
            peer_state.sync.reply_settled();
        }
        // While we catch up, we hold back sync with everyone but the source on purpose:
        // differing heads and idle sync are expected then, not a reason to heal.
        let catching_up = !state.initial_sync_done.load(Ordering::Acquire);
        !catching_up
            && peer_state.needs_heal(&local.heads, &remote.heads, &state.timing, Instant::now())
    };

    if heal {
        warn!(peer = %peer, "Documents differ and sync is idle; re-handshaking to restart sync");
        ensure_handshake(state, peer).await;
    }
}