//! Keepalives and the self-healing re-handshake.
//!
//! A quiet pair's keepalive (`PeerState::take_keepalive_due`) is a PING, answered with a
//! PONG by the peer's paths layer (`transport::paths`), over the same path: it tells both
//! sides the other is alive, and measures the path (`swarm::rtt`). Every `Timing::status_refresh` it is a
//! STATUS instead, answered with a STATUS_REPLY: each side's state (heads, readiness,
//! epoch), which also go out as they change (relay announcements, the epoch after a
//! switch, readiness: `announce`). When a STATUS shows our documents differ while sync
//! is idle, the pair re-handshakes (`heal_if_diverged`).

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

use anyhow::Result;
use tokio::time::Instant;
use tracing::{debug, info, 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 STATUS with our status.
pub(crate) async fn send_status(state: &AppState, peer: &str) -> Result<()> {
    let status = local_status(state, peer).await;
    if let Some(peer_state) = state.get_peer_state(peer).await {
        peer_state.lock().await.last_status_at = Some(Instant::now());
    }
    state
        .transport
        .send_message(peer, &ParsedMessage::Status(status))
        .await
}

/// Sends `peer` a PING (the paths layer measures the path with it).
pub(crate) async fn send_ping(state: &AppState, peer: &str) -> Result<()> {
    let nonce = crate::transport::paths::random_nonce();
    state
        .transport
        .send_message(peer, &ParsedMessage::Ping(nonce))
        .await
}

/// A PING or PONG from `peer`: it is alive, and after a PONG the paths layer has a
/// new round-trip time to it (`swarm::rtt`). A transport without the paths layer (`Transport::answers_pings`: in tests)
/// leaves answering a PING to us. Only peers we know already: a PING makes none.
pub(super) async fn note_alive(state: &AppState, peer: &str, message: &ParsedMessage) {
    let now = Instant::now();
    match message {
        ParsedMessage::Ping(nonce) if !state.transport.answers_pings() => {
            let pong = ParsedMessage::Pong(*nonce);
            if let Err(error) = state.transport.send_message(peer, &pong).await {
                debug!(peer = %peer, error = %format!("{error:#}"), "PONG failed");
            }
        }
        ParsedMessage::Pong(_) => {
            if let Some(rtt) = state.transport.rtt(peer) {
                state.lock_rtt().measured(peer, rtt, now);
            }
        }
        _ => {}
    }
    let Some(peer_state) = state.get_peer_state(peer).await else {
        return;
    };
    let reactivated = {
        let mut peer_state = peer_state.lock().await;
        if matches!(message, ParsedMessage::Ping(_)) {
            // Its keepalive (or a path probe): the next is ours.
            peer_state.keepalive_received_at = Some(now);
        }
        if peer_state.unreachable_since.is_some() {
            // Reachable again (`mark_seen` clears it).
            state.topology_changed.notify_one();
        }
        peer_state.mark_seen(now)
    };
    crate::swarm::source::heard_from(state, peer);
    if reactivated {
        info!(peer = %peer, "Dormant peer is active again; re-handshaking");
        ensure_handshake(state, peer).await;
    }
}

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

/// A STATUS to each neighbour, at most `within` (all sent: none).
async fn status_to_neighbours(state: &AppState, within: Option<Duration>) {
    let mut statuses = tokio::task::JoinSet::new();
    for (peer, _) in state.peers_snapshot().await {
        if topology::is_neighbour(state, &peer) {
            let state = state.clone();
            statuses.spawn(async move {
                if let Err(error) = send_status(&state, &peer).await {
                    debug!(peer = %peer, error = %format!("{error:#}"), "STATUS failed");
                }
            });
        }
    }
    let all_sent = async { while statuses.join_next().await.is_some() {} };
    match within {
        Some(within) => {
            let _ = tokio::time::timeout(within, all_sent).await;
        }
        None => all_sent.await,
    }
}

/// Our readiness changed (we became ready, or ready again after a rejoin): a STATUS
/// to each neighbour now, rather than with the next refresh. A peer opens sync towards
/// us only once it knows we are ready.
pub(crate) fn announce(state: &AppState) {
    let state = state.clone();
    tokio::spawn(async move { status_to_neighbours(&state, None).await });
}

/// At shutdown: one last STATUS 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);
    status_to_neighbours(state, Some(FAREWELL_TIMEOUT)).await;
}

/// Our status for `peer`.
pub(crate) async fn local_status(state: &AppState, peer: &str) -> transport::Status {
    // 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::Status {
        heads,
        epoch,
        epoch_seq,
        ready: state.is_ready().await && !state.stopping.load(Ordering::Acquire),
        syncs_with_you,
    }
}

/// A STATUS (not a reply) from `peer`: its keepalive, and a STATUS between us
/// (`Timing::status_refresh`).
pub(super) async fn status_received(state: &AppState, peer: &str) {
    if let Some(peer_state) = state.get_peer_state(peer).await {
        let mut peer_state = peer_state.lock().await;
        let now = Instant::now();
        peer_state.keepalive_received_at = Some(now);
        peer_state.last_status_at = Some(now);
    }
}

/// What a peer's status says about it.
pub(super) async fn note_status(state: &AppState, peer: &str, remote: &transport::Status) {
    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::Status,
    remote: &transport::Status,
) {
    // 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;
    }
}