//! 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, PeerState};
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 ping = matches!(message, ParsedMessage::Ping(_));
    let (reactivated, _) = heard_from(state, peer, &peer_state, now, |peer_state| {
        if ping {
            // Its keepalive (or a path probe): the next is ours.
            peer_state.keepalive_received_at = Some(now);
        }
    })
    .await;
    if reactivated {
        ensure_handshake(state, peer).await;
    }
}

/// A message from `peer` came in: it is seen (`PeerState::mark_seen`: active, and
/// reachable, again), and the catch-up source hears of it. `also` runs on its state
/// under the same lock. Returns whether it was dormant (the caller re-handshakes), and
/// when it was last seen before.
pub(super) async fn heard_from(
    state: &AppState,
    peer: &str,
    peer_state: &tokio::sync::Mutex<PeerState>,
    now: Instant,
    also: impl FnOnce(&mut PeerState),
) -> (bool, Option<Instant>) {
    let (reactivated, previously_seen) = {
        let mut peer_state = peer_state.lock().await;
        also(&mut peer_state);
        let previously_seen = peer_state.last_seen;
        if peer_state.unreachable_since.is_some() {
            // Reachable again (`mark_seen` clears it).
            state.topology_changed.notify_one();
        }
        (peer_state.mark_seen(now), previously_seen)
    };
    crate::swarm::source::heard_from(state, peer);
    if reactivated {
        info!(peer = %peer, "Dormant peer is active again; re-handshaking");
    }
    (reactivated, previously_seen)
}

/// 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.closing.send_replace(true);
    info!("Saying farewell to our neighbours");
    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) => peer_state.lock().await.syncs_on(epoch),
        None => false,
    };
    transport::Status {
        heads,
        epoch,
        epoch_seq,
        ready: state.is_ready().await && !state.is_closing(),
        syncs_with_you,
        leaving: state.is_closing(),
    }
}

/// 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;
        // Its farewell: dormant now, as if silent for long (a message from it, as its
        // HELLO when it's back, makes it active again).
        let left = remote.leaving && peer_state.say_farewell(Instant::now());
        let became_ready = remote.ready && peer_state.reports_ready != Some(true);
        peer_state.reports_ready = Some(remote.ready);
        let changed =
            peer_state.syncs_with_us.replace(remote.syncs_with_you) != Some(remote.syncs_with_you);
        peer_state.reported_heads = remote.heads.heads().map(<[_]>::len);
        drop(peer_state);
        if left {
            info!(peer = %peer, "Peer said farewell; dormant until it is heard from again");
        }
        if changed || left {
            // A neighbour that only connects out stays one as long as it says so.
            state.topology_changed.notify_one();
        }
        // 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.
            state.request_sync(peer);
        }
    }
    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;
    }
}