//! 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 active peer advertises a newer generation, the node downloads that peer's
//! document (`SnapshotRequest`/`SnapshotReply`), checks that it 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).
//!
//! The document travels in chunks of `CHUNK_BYTES`, one request per chunk, so a large
//! document neither needs one huge transport message nor one long timeout. The serving
//! node saves its document once per download and serves every chunk from that copy; a
//! SHA-256 digest of the whole copy is in every reply and checked at the end.
//!
//! Peers refuse explicitly (non-members, busy, other epoch), so the requester moves on
//! to the next peer right away instead of timing out.

use std::sync::Arc;
use std::time::Duration;

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

use super::delta::{self, ReplayStats};
use super::local::{CatchUp, LocalStatus, Served};
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};

/// Pause before the next attempt (next peer) after a failed download.
const RETRY: Duration = Duration::from_secs(10);
/// Payload per reply.
pub(crate) const CHUNK_BYTES: usize = 1024 * 1024;
/// How long to wait for one chunk, and how often to ask for it.
const CHUNK_TIMEOUT: Duration = Duration::from_secs(60);
const CHUNK_ATTEMPTS: u32 = 3;
/// A served copy older than this is replaced when a new download starts.
const SERVED_FRESH_FOR: 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();

    // Any active peer will do: it proved liveness recently, and whether it accepts us
    // as a member is its decision (it refuses explicitly otherwise).
    let mut ahead: Vec<(String, Epoch)> = host
        .peer_summaries()
        .await
        .into_iter()
        .filter(|peer| peer.active)
        .filter_map(|peer| peer.epoch.map(|epoch| (peer.id, epoch)))
        .filter(|(_, epoch)| epoch.generation > local.generation)
        .collect();
    ahead.sort_by(|a, b| a.0.cmp(&b.0));

    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 (sorted, so the rotation is stable).
    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 download(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),
            }));
        }
    }
}

/// Downloads `peer`'s document of `target`, chunk by chunk.
async fn download<H: Host>(
    host: &H,
    peer: &str,
    local: Epoch,
    target: Epoch,
) -> Result<Vec<u8>, String> {
    let mut document: Vec<u8> = Vec::new();
    let mut expected: Option<([u8; 32], u64)> = None;
    let mut last_progress = 0;

    loop {
        let offset = document.len() as u64;
        let reply = request_chunk(host, peer, local, target, offset).await?;

        if !reply.refused.is_empty() {
            return Err(format!("peer refused: {}", reply.refused));
        }
        if reply.epoch != target {
            return Err(format!("peer is on {} instead of {target}", reply.epoch));
        }
        match expected {
            None => {
                info!(peer = %peer, document_bytes = reply.total, "Downloading document");
                expected = Some((reply.digest, reply.total));
            }
            Some((digest, total)) if digest != reply.digest || total != reply.total => {
                return Err("the peer's document changed during the download".to_owned());
            }
            Some(_) => {}
        }
        let total = reply.total;
        if reply.data.is_empty() && offset < total {
            return Err(format!("empty chunk at offset {offset} of {total}"));
        }
        document.extend_from_slice(&reply.data);

        let percent = (document.len() as u64 * 100)
            .checked_div(total)
            .unwrap_or(100);
        if percent >= last_progress + 10 {
            last_progress = percent;
            debug!(peer = %peer, percent, "Download progress");
        }

        if document.len() as u64 >= total {
            break;
        }
    }

    let (digest, total) = expected.expect("set by the first reply");
    if document.len() as u64 != total || sha256(&document) != digest {
        return Err("downloaded document does not match its digest".to_owned());
    }
    Ok(document)
}

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

        match tokio::time::timeout(CHUNK_TIMEOUT, reply).await {
            Ok(Ok(reply)) => return Ok(reply),
            Ok(Err(_)) => return Err("request superseded".to_owned()),
            Err(_) => {
                host.compaction().stop_awaiting();
                debug!(peer = %peer, offset, attempt, "No reply for chunk; asking again");
            }
        }
    }
    Err(format!(
        "no reply for the chunk at offset {offset} after {CHUNK_ATTEMPTS} attempts of {}s",
        CHUNK_TIMEOUT.as_secs()
    ))
}

/// 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 chunk request from a saved copy of this node's document, if this
/// node is on the wanted epoch and the peer is a member of our cluster. Refuses
/// explicitly otherwise.
pub(crate) async fn serve_snapshot<H: Host>(host: H, peer: String, request: SnapshotRequest) {
    let current = host.current_epoch();
    let wanted = request.wanted;

    let reply = match served_copy(&host, &peer, current, wanted, request.offset).await {
        Ok(served) => {
            let start = (request.offset as usize).min(served.bytes.len());
            let end = (start + CHUNK_BYTES).min(served.bytes.len());
            if request.offset == 0 {
                info!(peer = %peer, wanted = %wanted, document_bytes = served.bytes.len(), "Serving document");
            }
            SnapshotReply {
                epoch: served.epoch,
                wanted,
                digest: served.digest,
                total: served.bytes.len() as u64,
                offset: request.offset,
                data: served.bytes[start..end].to_vec(),
                refused: String::new(),
            }
        }
        Err(reason) => {
            info!(peer = %peer, wanted = %wanted, reason = %reason, "Refusing document request");
            SnapshotReply {
                epoch: current,
                wanted,
                digest: [0; 32],
                total: 0,
                offset: request.offset,
                data: Vec::new(),
                refused: reason,
            }
        }
    };

    if let Err(error) = host
        .send_message(&peer, ParsedMessage::SnapshotReply(reply))
        .await
    {
        warn!(peer = %peer, error = %format!("{error:#}"), "Sending document chunk failed");
    }
}

/// The saved document to serve chunks from: reused while a download is in progress,
/// saved anew when a download starts and the copy is stale.
async fn served_copy<H: Host>(
    host: &H,
    peer: &str,
    current: Epoch,
    wanted: Epoch,
    offset: u64,
) -> Result<Arc<Served>, String> {
    if !host.is_cluster_member(peer).await {
        return Err("not a member of this cluster (no completed handshake)".to_owned());
    }
    if current != wanted {
        return Err(format!("this node is on {current}"));
    }

    if let Some(served) = host.compaction().served(wanted) {
        if offset > 0 || served.saved_at.elapsed() < SERVED_FRESH_FOR {
            return Ok(served);
        }
    } else if offset > 0 {
        return Err("download expired; start again".to_owned());
    }

    if !host.compaction().try_start_serving() {
        return Err("busy saving the document; try again".to_owned());
    }

    // 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;
    let saved = tokio::task::spawn_blocking(move || {
        if doc
            .get_change_by_hash(&automerge::ChangeHash(wanted.id))
            .is_some()
        {
            Some(doc.save())
        } else {
            None
        }
    })
    .await;
    host.compaction().done_serving();

    let bytes = match saved {
        Ok(Some(bytes)) => bytes,
        Ok(None) => return Err("this node switched epochs".to_owned()),
        Err(error) => {
            error!(error = %error, "Saving the document for a peer failed");
            return Err("saving the document failed".to_owned());
        }
    };

    let served = Arc::new(Served {
        epoch: wanted,
        digest: sha256(&bytes),
        bytes,
        saved_at: Instant::now(),
    });
    host.compaction().set_served(Arc::clone(&served));
    Ok(served)
}

fn sha256(bytes: &[u8]) -> [u8; 32] {
    let mut digest = [0; 32];
    digest.copy_from_slice(&Sha256::digest(bytes));
    digest
}

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