//! Catching up with a newer epoch: the node rejoins.
//!
//! Peers on another epoch don't sync, so a node that missed a switch (new, late,
//! returning, or one whose own build failed) can't follow through ordinary sync. Once
//! an active peer advertises a newer generation, the node
//!
//! 1. stops being ready and keeps its document aside; reads are answered from a copy
//!    of it, marked stale (`X-Continuum-Stale`), writes get `503`;
//! 2. switches to an empty document labelled with the newer epoch: an empty history
//!    conflicts with nothing, so ordinary sync fills it from peers on that epoch,
//!    exactly like a new node joining;
//! 3. once caught up with a peer, checks that the history starts with the epoch's
//!    snapshot change;
//! 4. decides about its own changes the swarm never saw (none, for a new node):
//!    * the metadata of the epoch that followed the node's own records which changes
//!      went into it (per-actor clock); the rest are this node's missing changes;
//!    * within that epoch's replay window they are replayed (state per path, as in a
//!      switch), otherwise they are quarantined for an operator
//!      (`/admin/quarantine`), as are changes that can't be told apart;
//! 5. archives the old document and serves requests again. A node that was never
//!    ready (a new node) is made ready by bootstrap, which waits for the rejoin.
//!
//! If no peer completes a sync in time, or the check fails, the node restores its old
//! document and epoch and tries again later.

use std::collections::BTreeMap;
use std::convert::Infallible;
use std::time::Duration;

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

use super::delta::{self, ReplayStats};
use super::kind::Translator;
use super::local::{CatchUp, LocalStatus, Rejoin};
use super::quarantine;
use super::records::unix_millis;
use super::snapshot::{self, EpochMetadata};
use super::Host;
use crate::document::replica::Epoch;

/// Pause before trying the same target again after a failed rejoin.
const RETRY: Duration = Duration::from_secs(30);
/// Lower bound for how long a rejoin may take to complete a sync.
const MIN_REJOIN_TIMEOUT: Duration = Duration::from_secs(60);

/// Called by the switch worker when there is nothing to switch to locally, and on every
/// tick while a rejoin is in progress.
pub(super) async fn check<H: Host>(host: &H) {
    if let Some((target, started)) = host.remake().rejoin_progress() {
        let timeout = host.dormant_after().max(MIN_REJOIN_TIMEOUT);
        if host.initial_sync_done() {
            finish(host).await;
        } else if started.elapsed() > timeout {
            if let Some(rejoin) = host.remake().take_rejoin() {
                let reason = format!(
                    "no peer on {target} completed a sync within {}s",
                    timeout.as_secs()
                );
                abandon(host, rejoin, reason).await;
            }
        }
        return;
    }

    let local = host.current_epoch();
    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.remake().catch_up().is_some() {
            host.remake().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
        .remake()
        .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;

    start(host, local, target, attempts).await;
}

async fn start<H: Host>(host: &H, local: Epoch, target: Epoch, attempts: u32) {
    host.remake().set_catch_up(Some(CatchUp {
        target,
        attempts,
        last_attempt: Instant::now(),
        error: None,
    }));

    // Keep answering reads, from a copy of the current document marked stale, before
    // the node stops being ready: no window in which reads get 503.
    if host.is_node_ready().await {
        let copy = host.read_doc(|doc| doc.clone()).await;
        host.set_stale_view(Some((local, copy)));
    }
    let was_ready = host.suspend().await;
    let old = match host
        .switch_epoch(target, |_live| Ok::<_, Infallible>(AutoCommit::new()))
        .await
    {
        Ok(old) => old,
        Err(never) => match never {},
    };
    // After the swap: messages of the old epoch are dropped from now on, so nothing
    // can mark the new, empty document as caught up by mistake. A stale "done" from
    // just before is caught in `finish` (the snapshot change is missing then).
    host.reset_initial_sync();

    info!(from = %local, to = %target, attempt = attempts, "Rejoining newer epoch");
    host.remake().start_rejoin(Rejoin {
        previous: local,
        target,
        old,
        started: Instant::now(),
        was_ready,
    });
}

/// What to do with this node's own changes the new epoch lacks.
enum Plan {
    Nothing,
    /// Replay these paths, carried through every remake since the node's epoch.
    Replay(Vec<Vec<String>>, Translator),
    Quarantine {
        paths: Vec<Vec<String>>,
        reason: String,
    },
}

enum Prepared {
    /// Caught up with someone, but not with the target epoch yet: keep waiting.
    NotYet(AutoCommit),
    Failed(AutoCommit, String),
    Ready {
        old: AutoCommit,
        metadata: EpochMetadata,
        plan: Plan,
        own_missing: usize,
        had_history: bool,
    },
}

async fn finish<H: Host>(host: &H) {
    let Some(rejoin) = host.remake().take_rejoin() else {
        return;
    };
    let Rejoin {
        previous,
        target,
        old,
        started,
        was_ready,
    } = rejoin;

    // Everything heavy runs on a copy, without a lock.
    let copy = host.read_doc(|doc| doc.clone()).await;
    let prepared = tokio::task::spawn_blocking(move || prepare(copy, old, previous, target)).await;

    let (old, metadata, plan, own_missing, had_history) = match prepared {
        Ok(Prepared::Ready {
            old,
            metadata,
            plan,
            own_missing,
            had_history,
        }) => (old, metadata, plan, own_missing, had_history),
        Ok(Prepared::NotYet(old)) => {
            host.reset_initial_sync();
            host.remake().start_rejoin(Rejoin {
                previous,
                target,
                old,
                started,
                was_ready,
            });
            return;
        }
        Ok(Prepared::Failed(old, reason)) => {
            let rejoin = Rejoin {
                previous,
                target,
                old,
                started,
                was_ready,
            };
            return abandon(host, rejoin, reason).await;
        }
        Err(join_error) => {
            // The old document is gone with the panicked task; nothing to restore.
            error!(error = %join_error, "Rejoin task failed; staying on the new epoch");
            if was_ready {
                host.resume().await;
            }
            host.set_stale_view(None);
            return;
        }
    };

    let mut stats = ReplayStats::default();
    let mut quarantine_id = None;
    let mut keep_archive = true;
    let mut old = Some(old);

    match plan {
        Plan::Nothing => {}
        Plan::Replay(paths, translator) => {
            let document = old.take().expect("old document");
            let (replayed, document) = host
                .commit_doc(move |doc| {
                    let stats = delta::replay_paths(&document, doc, &paths, &translator);
                    Ok::<_, Infallible>((stats, document))
                })
                .await
                .unwrap_or_else(|never| match never {});
            stats = replayed;
            // Writes the remakes since had no place for.
            quarantine_id = quarantine::add(
                host,
                previous,
                target,
                &document,
                std::mem::take(&mut stats.quarantined),
            );
            old = Some(document);
        }
        Plan::Quarantine { paths, reason } => {
            let document = old.take().expect("old document");
            let entries = paths
                .into_iter()
                .map(|path| (path, reason.clone()))
                .collect();
            quarantine_id = quarantine::add(host, previous, target, &document, entries);
            keep_archive = false; // the quarantine holds the document
        }
    }

    if let (true, true, Some(document)) = (had_history, keep_archive, old) {
        host.remake().archive(
            previous,
            document,
            metadata.policy.archive_retention.clone(),
            target.generation,
        );
    }

    host.remake().set_status(
        &metadata.remake_id,
        LocalStatus {
            switched_at: Some(unix_millis()),
            replayed_paths: Some(stats.replayed),
            failed_paths: Some(stats.failed),
            quarantine: quarantine_id.clone(),
            error: None,
        },
    );
    host.remake().set_catch_up(None);
    if was_ready {
        host.resume().await;
    }
    // After resuming, so reads move from the stale view to the live document directly.
    host.set_stale_view(None);

    info!(
        from = %previous,
        to = %target,
        own_changes = own_missing,
        replayed = stats.replayed,
        failed = stats.failed,
        quarantine = ?quarantine_id,
        secs = started.elapsed().as_secs(),
        "Rejoined newer epoch"
    );
}

/// Verifies the synced document and decides about the old document's own changes.
fn prepare(mut copy: AutoCommit, mut old: AutoCommit, previous: Epoch, target: Epoch) -> Prepared {
    if copy.get_change_by_hash(&ChangeHash(target.id)).is_none() {
        return Prepared::NotYet(old);
    }
    let metadata = match snapshot::verify_document(&mut copy, target) {
        Ok(metadata) => metadata,
        Err(reason) => return Prepared::Failed(old, reason),
    };

    let old_heads = old.get_heads();
    if old_heads.is_empty() {
        return Prepared::Ready {
            old,
            metadata,
            plan: Plan::Nothing,
            own_missing: 0,
            had_history: false,
        };
    }

    // The epoch that followed ours records (per actor) which of our changes it took.
    // Snapshots carry earlier epochs' metadata forward, so this works across any gap.
    let next = snapshot::read_epoch_metadata(&copy, previous.generation + 1)
        .ok()
        .filter(|next| next.previous == previous);

    let own = old.get_actor().clone();
    let clock = next
        .as_ref()
        .map_or_else(BTreeMap::new, |next| next.clock.clone());
    let missing = delta::missing_from_snapshot(&mut old, &clock, &own);
    if missing.own == 0 {
        return Prepared::Ready {
            old,
            metadata,
            plan: Plan::Nothing,
            own_missing: 0,
            had_history: true,
        };
    }

    let paths = delta::changed_paths(&mut old, &missing.before, &old_heads);
    let gap = target.generation - previous.generation;

    // Writes made in our epoch go through every remake since, in order.
    let translator = (previous.generation + 1..=target.generation).try_fold(
        Translator::default(),
        |mut translator, generation| {
            let epoch = snapshot::read_epoch_metadata(&copy, generation)?;
            let kind = epoch.kind().ok_or_else(|| {
                format!(
                    "epoch generation {generation} was made by unknown remake kind {:?}",
                    epoch.kind
                )
            })?;
            translator.push(kind, epoch.translation);
            Ok::<_, String>(translator)
        },
    );

    let plan = match (&next, translator) {
        (Some(_), Err(reason)) => Plan::Quarantine {
            paths,
            reason: format!("cannot translate {} local changes: {reason}", missing.own),
        },
        (None, _) => Plan::Quarantine {
            paths,
            reason: format!(
                "the metadata of the epoch after {previous} is missing, so it is unknown which of {} local changes the swarm has",
                missing.own
            ),
        },
        (Some(next), Ok(translator)) => {
            let away_ms = unix_millis().saturating_sub(next.proposed_at);
            if next.policy.replay_window.covers(gap, away_ms) {
                Plan::Replay(paths, translator)
            } else {
                Plan::Quarantine {
                    paths,
                    reason: format!(
                        "{} local changes are outside the replay window ({gap} epochs behind, swarm moved on {}s ago)",
                        missing.own,
                        away_ms / 1000
                    ),
                }
            }
        }
    };

    Prepared::Ready {
        old,
        metadata,
        plan,
        own_missing: missing.own,
        had_history: true,
    }
}

/// Restores the old document and epoch, and retries later.
async fn abandon<H: Host>(host: &H, rejoin: Rejoin, reason: String) {
    let Rejoin {
        previous,
        target,
        old,
        was_ready,
        ..
    } = rejoin;
    warn!(target = %target, reason = %reason, "Rejoin failed; restoring the previous epoch");

    if let Err(never) = host
        .switch_epoch(previous, move |_live| Ok::<_, Infallible>(old))
        .await
    {
        match never {}
    }
    // A node that was never ready must sync again before bootstrap may finish.
    host.reset_initial_sync();
    if was_ready {
        host.resume().await;
    }
    host.set_stale_view(None);

    let attempts = host
        .remake()
        .catch_up()
        .filter(|catch_up| catch_up.target == target)
        .map_or(1, |catch_up| catch_up.attempts);
    host.remake().set_catch_up(Some(CatchUp {
        target,
        attempts,
        last_attempt: Instant::now(),
        error: Some(reason),
    }));
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::remake::snapshot::{build, SnapshotParams};
    use automerge::{transaction::Transactable, ActorId, ROOT};
    use serde_json::{json, Value as JsonValue};

    /// An old-epoch document with one own change after the base, and the swarm's
    /// document of the next epoch (snapshot plus one foreign change).
    fn scenario(policy: &JsonValue) -> (AutoCommit, AutoCommit, Epoch) {
        let mut old = AutoCommit::new().with_actor(ActorId::from(&b"late-node"[..]));
        old.put(ROOT, "base", 1_i64).unwrap();
        old.commit();
        let base = old.get_heads();

        let snapshot = build(
            &mut old,
            &SnapshotParams {
                remake_id: "c-1",
                previous: Epoch::GENESIS,
                base_heads: &base,
                policy,
                proposed_at: 0,
                kind: "compact",
                content: &JsonValue::Null,
                translation: &JsonValue::Null,
            },
        )
        .unwrap();
        let target = snapshot.epoch;

        let mut swarm = snapshot.doc.with_actor(ActorId::from(&b"swarm-node"[..]));
        swarm.put(ROOT, "theirs", 1_i64).unwrap();
        swarm.commit();

        old.put(ROOT, "mine", 1_i64).unwrap();
        old.commit();
        (old, swarm, target)
    }

    #[test]
    fn replays_within_the_window() {
        let (old, swarm, target) = scenario(&json!({}));
        match prepare(swarm, old, Epoch::GENESIS, target) {
            Prepared::Ready {
                plan: Plan::Replay(paths, _),
                own_missing,
                ..
            } => {
                assert_eq!(own_missing, 1);
                assert_eq!(paths, vec![vec!["mine".to_owned()]]);
            }
            _ => panic!("expected a replay"),
        }
    }

    #[test]
    fn quarantines_outside_the_window() {
        let (old, swarm, target) = scenario(&json!({ "replayWindow": { "epochs": 0 } }));
        match prepare(swarm, old, Epoch::GENESIS, target) {
            Prepared::Ready {
                plan: Plan::Quarantine { paths, .. },
                ..
            } => assert_eq!(paths, vec![vec!["mine".to_owned()]]),
            _ => panic!("expected a quarantine"),
        }
    }

    #[test]
    fn new_node_has_nothing_to_replay() {
        let (_, swarm, target) = scenario(&json!({}));
        assert!(matches!(
            prepare(swarm, AutoCommit::new(), Epoch::GENESIS, target),
            Prepared::Ready {
                plan: Plan::Nothing,
                had_history: false,
                ..
            }
        ));
    }

    #[test]
    fn waits_for_the_target_epoch() {
        let (old, _, target) = scenario(&json!({}));
        assert!(matches!(
            prepare(AutoCommit::new(), old, Epoch::GENESIS, target),
            Prepared::NotYet(_)
        ));
    }
}