//! Catching up with a newer epoch from a peer.
//!
//! Peers on another epoch don't sync, so a node that missed a switch (new, late, or
//! one whose own build failed) can't learn the new epoch through ordinary sync. Instead,
//! once an established peer advertises a newer generation, the node asks it for its
//! whole document (`SnapshotRequest`), checks that the document starts with the
//! advertised snapshot change, and switches to it:
//!
//! * a node that wrote nothing the snapshot lacks (e.g. a brand-new node) adopts the
//!   document as is, whatever the gap;
//! * a node exactly one epoch behind replays its own changes the snapshot lacks, found
//!   with the snapshot's per-actor clock;
//! * a node further behind with such changes stays where it is and reports it; that is
//!   the returning-node case (replay window or quarantine, step 7).
//!
//! Only one fetch runs at a time, and the worker retries after `RETRY`. The reply is the
//! peer's full document in one message; chunking comes later.

use std::time::Duration;

use automerge::{ActorId, AutoCommit};
use tokio::time::Instant;
use tracing::{debug, error, info, warn};

use super::delta::{self, ReplayStats};
use super::local::{CatchUp, LocalStatus};
use super::records::unix_millis;
use super::snapshot;
use super::switch::swap_in;
use super::Host;
use crate::replica::Epoch;
use crate::transport::{ParsedMessage, SnapshotReply, SnapshotRequest};

const RETRY: Duration = Duration::from_secs(10);
const FETCH_TIMEOUT: Duration = Duration::from_secs(60);

/// Called by the switch worker when there is nothing to switch to locally.
pub(super) async fn check<H: Host>(host: &H) {
    let local = host.current_epoch();

    let ahead: Vec<(String, Epoch)> = host
        .peer_summaries()
        .await
        .into_iter()
        .filter(|peer| peer.active && peer.established)
        .filter_map(|peer| peer.epoch.map(|epoch| (peer.id, epoch)))
        .filter(|(_, epoch)| epoch.generation > local.generation)
        .collect();

    let Some(generation) = ahead.iter().map(|(_, epoch)| epoch.generation).max() else {
        if host.compaction().catch_up().is_some() {
            host.compaction().set_catch_up(None);
        }
        return;
    };

    let mut targets: Vec<Epoch> = ahead
        .iter()
        .filter(|(_, epoch)| epoch.generation == generation)
        .map(|(_, epoch)| *epoch)
        .collect();
    targets.sort_by_key(|epoch| epoch.id);
    targets.dedup();
    let [target] = targets.as_slice() else {
        let listed: Vec<String> = targets.iter().map(ToString::to_string).collect();
        warn!(epochs = ?listed, "Peers disagree about the newest epoch; not catching up");
        return;
    };
    let target = *target;

    // Back off between attempts for the same target.
    let previous = host
        .compaction()
        .catch_up()
        .filter(|catch_up| catch_up.target == target);
    if previous
        .as_ref()
        .is_some_and(|catch_up| catch_up.last_attempt.elapsed() < RETRY)
    {
        return;
    }
    let attempts = previous.map_or(0, |catch_up| catch_up.attempts) + 1;

    // Rotate through the peers on the target epoch.
    let sources: Vec<&String> = ahead
        .iter()
        .filter(|(_, epoch)| *epoch == target)
        .map(|(peer, _)| peer)
        .collect();
    let source = sources[(attempts as usize - 1) % sources.len()].clone();

    host.compaction().set_catch_up(Some(CatchUp {
        target,
        attempts,
        last_attempt: Instant::now(),
        error: None,
    }));

    info!(local = %local, target = %target, peer = %source, attempt = attempts, "Fetching newer epoch from peer");

    let result = match fetch(host, &source, local, target).await {
        Ok(document) => adopt(host, local, target, document).await,
        Err(error) => Err(error),
    };

    match result {
        Ok(()) => host.compaction().set_catch_up(None),
        Err(error) => {
            warn!(target = %target, peer = %source, error = %error, "Catching up failed; will retry");
            host.compaction().set_catch_up(Some(CatchUp {
                target,
                attempts,
                last_attempt: Instant::now(),
                error: Some(error),
            }));
        }
    }
}

async fn fetch<H: Host>(
    host: &H,
    peer: &str,
    local: Epoch,
    target: Epoch,
) -> Result<Vec<u8>, String> {
    let reply = host.compaction().await_reply(peer, target);
    let request = ParsedMessage::SnapshotRequest(SnapshotRequest {
        epoch: local,
        wanted: target,
    });
    if let Err(error) = host.send_message(peer, request).await {
        host.compaction().stop_awaiting();
        return Err(format!("sending the request failed: {error:#}"));
    }

    let reply = match tokio::time::timeout(FETCH_TIMEOUT, reply).await {
        Ok(Ok(reply)) => reply,
        Ok(Err(_)) => return Err("request superseded".to_owned()),
        Err(_) => {
            host.compaction().stop_awaiting();
            return Err(format!("no reply within {}s", FETCH_TIMEOUT.as_secs()));
        }
    };

    if reply.epoch != target || reply.document.is_empty() {
        return Err(format!(
            "peer is on {} and did not send {target}",
            reply.epoch
        ));
    }
    Ok(reply.document)
}

/// Verifies the fetched document, carries over this node's own missing changes, and
/// switches to it.
async fn adopt<H: Host>(
    host: &H,
    local: Epoch,
    target: Epoch,
    document: Vec<u8>,
) -> Result<(), String> {
    // Copy under the read lock only; everything else happens without a lock.
    let copy = host.read_doc(|doc| doc.clone()).await;
    let bytes = document.len();

    let prepared = tokio::task::spawn_blocking(move || -> Result<Prepared, String> {
        let mut fetched =
            AutoCommit::load(&document).map_err(|error| format!("invalid document: {error}"))?;
        let metadata = snapshot::verify_document(&mut fetched, target)?;

        let mut copy = copy;
        let own = copy.get_actor().clone();
        let copy_heads = copy.get_heads();
        let had_history = !copy_heads.is_empty();
        let missing = delta::missing_from_snapshot(&mut copy, &metadata.clock, &own);

        let mut doc = fetched.with_actor(ActorId::random());
        let stats = if missing.own == 0 {
            ReplayStats::default()
        } else if metadata.previous == local {
            delta::replay(&mut copy, &missing.before, &copy_heads, &mut doc)
        } else {
            return Err(format!(
                "{} local changes are not in {target}, and this node is more than one epoch behind (on {local}, {target} follows {}); not replaying them automatically",
                missing.own, metadata.previous
            ));
        };

        Ok(Prepared {
            doc,
            copy_heads,
            stats,
            own_missing: missing.own,
            had_history,
            metadata,
        })
    })
    .await
    .map_err(|error| format!("catch-up task failed: {error}"))??;

    let Prepared {
        doc,
        copy_heads,
        stats,
        own_missing,
        had_history,
        metadata,
    } = prepared;

    let old = swap_in(host, target, doc, copy_heads).await?;

    if had_history {
        host.compaction().archive(
            local,
            old,
            metadata.policy.archive_retention.clone(),
            target.generation,
        );
    }

    // Visible in the compaction's report, if this node has its record.
    host.compaction().set_status(
        &metadata.compaction_id,
        LocalStatus {
            switched_at: Some(unix_millis()),
            replayed_paths: Some(stats.replayed),
            failed_paths: Some(stats.failed),
            error: None,
        },
    );

    info!(
        from = %local,
        to = %target,
        document_bytes = bytes,
        own_changes = own_missing,
        replayed = stats.replayed,
        failed = stats.failed,
        "Caught up with newer epoch"
    );
    Ok(())
}

struct Prepared {
    doc: AutoCommit,
    copy_heads: Vec<automerge::ChangeHash>,
    stats: ReplayStats,
    own_missing: usize,
    had_history: bool,
    metadata: snapshot::EpochMetadata,
}

/// Answers a peer's snapshot request with this node's document, if it is on the wanted
/// epoch. Only for members of our cluster; one at a time.
pub(crate) async fn serve_snapshot<H: Host>(host: H, peer: String, request: SnapshotRequest) {
    if !host.is_cluster_member(&peer).await {
        debug!(peer = %peer, "Ignoring snapshot request from a non-member");
        return;
    }
    if !host.compaction().try_start_serving() {
        debug!(peer = %peer, "Already serving a snapshot; the peer will retry");
        return;
    }

    let current = host.current_epoch();
    let wanted = request.wanted;

    let document = if current == wanted {
        // Clone under the read lock, save without one. The history check guards against
        // a switch between reading the epoch and cloning.
        let mut doc = host.read_doc(|doc| doc.clone()).await;
        tokio::task::spawn_blocking(move || {
            if doc
                .get_change_by_hash(&automerge::ChangeHash(wanted.id))
                .is_some()
            {
                doc.save()
            } else {
                Vec::new()
            }
        })
        .await
        .unwrap_or_else(|error| {
            error!(error = %error, "Saving the document for a peer failed");
            Vec::new()
        })
    } else {
        Vec::new()
    };
    host.compaction().done_serving();

    info!(peer = %peer, wanted = %wanted, document_bytes = document.len(), "Serving snapshot");
    let reply = ParsedMessage::SnapshotReply(SnapshotReply {
        epoch: if document.is_empty() { current } else { wanted },
        document,
    });
    if let Err(error) = host.send_message(&peer, reply).await {
        warn!(peer = %peer, error = %format!("{error:#}"), "Sending snapshot failed");
    }
}

/// Hands a peer's snapshot reply to the waiting fetch.
pub(crate) fn receive_snapshot_reply<H: Host>(host: &H, peer: &str, reply: SnapshotReply) {
    if !host.compaction().deliver_reply(peer, reply) {
        debug!(peer = %peer, "Unexpected snapshot reply; dropped");
    }
}