//! Replaying a node's changes after the base into the new epoch.
//!
//! Changes from the old epoch can't be applied to the new document (their history is
//! unrelated), so they are replayed as state: for every path that changed between two
//! sets of heads, the new document gets whatever the old document's merged view holds
//! there now (or a delete if nothing). Replayed values are fresh operations, so they win
//! over what the snapshot holds at the same path.
//!
//! Granularity:
//! * Map entries are replayed per key, so changes to sibling keys made on different
//!   nodes don't overwrite each other.
//! * Lists and text are replayed whole: concurrent edits of the same list or text on two
//!   nodes around a compaction resolve as last-writer-wins.
//! * Edits inside a map that lost a merge conflict produce no patch in Automerge and are
//!   not replayed. Such maps only exist briefly between concurrent creations.
//!
//! The protocol's own metadata (`_continuum/compactions`, `_continuum/epochs`) is never
//! replayed: the snapshot owns it.

use std::collections::{BTreeMap, HashSet};

use automerge::{ActorId, AutoCommit, ChangeHash, PatchAction, Prop};
use tracing::warn;

use super::snapshot::{COMPACTIONS, EPOCHS, META_ROOT};
use crate::store::{self, StoreError};

#[derive(Debug, Default, Clone, Copy)]
pub(crate) struct ReplayStats {
    pub(crate) replayed: usize,
    pub(crate) failed: usize,
}

/// Replays into `dst` everything that changed in `src` between `before` and `after`,
/// as one commit. Both head sets must be in `src`.
pub(crate) fn replay(
    src: &mut AutoCommit,
    before: &[ChangeHash],
    after: &[ChangeHash],
    dst: &mut AutoCommit,
) -> ReplayStats {
    if before == after {
        return ReplayStats::default();
    }

    let paths = changed_paths(src, before, after);
    replay_paths(src, dst, &paths)
}

/// Replays the state of `paths` (from [`changed_paths`]) from `src` into `dst`, as one
/// commit. Lets callers find the paths outside a lock and copy them inside it.
pub(crate) fn replay_paths(
    src: &AutoCommit,
    dst: &mut AutoCommit,
    paths: &[Vec<String>],
) -> ReplayStats {
    let mut stats = ReplayStats::default();
    for path in paths {
        replay_path(src, dst, path, &mut stats);
    }
    let _ = dst.commit();
    stats
}

fn replay_path(src: &AutoCommit, dst: &mut AutoCommit, path: &[String], stats: &mut ReplayStats) {
    if is_protocol_path(path) {
        return;
    }

    // Replacing an ancestor of the protocol metadata as a whole would clobber it; replay
    // its children one by one instead.
    if contains_protocol_path(path) {
        let mut keys = store::keys_at(src, path).unwrap_or_default();
        keys.extend(store::keys_at(dst, path).unwrap_or_default());
        keys.sort();
        keys.dedup();
        for key in keys {
            let mut child = path.to_vec();
            child.push(key);
            replay_path(src, dst, &child, stats);
        }
        return;
    }

    match store::copy_path(src, dst, path) {
        Ok(()) => stats.replayed += 1,
        Err(error) => {
            log_failure(path, &error);
            stats.failed += 1;
        }
    }
}

fn log_failure(path: &[String], error: &StoreError) {
    warn!(path = ?path, error = %error, "Could not replay change into the new epoch");
}

fn is_protocol_path(path: &[String]) -> bool {
    matches!(path, [root, part, ..] if root == META_ROOT && (part == COMPACTIONS || part == EPOCHS))
}

/// True for the root and for `_continuum` itself.
fn contains_protocol_path(path: &[String]) -> bool {
    matches!(path, [] | [_]) && path.first().is_none_or(|root| root == META_ROOT)
}

/// What a late node has that a snapshot lacks.
#[derive(Debug)]
pub(crate) struct Missing {
    /// Heads of everything except `own` changes the snapshot lacks (and changes built on
    /// them). Replaying from these heads to the document's heads replays exactly those.
    pub(crate) before: Vec<ChangeHash>,
    /// Changes by `own` that the snapshot lacks.
    pub(crate) own: usize,
}

/// Splits `doc`'s history against a snapshot `clock`: a change (actor, seq) is in the
/// snapshot iff seq <= clock[actor].
///
/// Only changes this node wrote itself (actor `own`) are worth replaying: other nodes'
/// changes after the base were replayed by those nodes when they switched, and replaying
/// them again would overwrite newer writes in the new epoch with stale ones.
pub(crate) fn missing_from_snapshot(
    doc: &mut AutoCommit,
    clock: &BTreeMap<String, u64>,
    own: &ActorId,
) -> Missing {
    let own_hex = own.to_string();
    let changes: Vec<(ChangeHash, Vec<ChangeHash>, bool)> = doc
        .get_changes(&[])
        .iter()
        .map(|change| {
            let actor = change.actor_id().to_string();
            let in_snapshot = clock.get(&actor).is_some_and(|&seq| change.seq() <= seq);
            let own_missing = actor == own_hex && !in_snapshot;
            (change.hash(), change.deps().to_vec(), own_missing)
        })
        .collect();

    // Excluded: own missing changes and everything that (transitively) depends on them.
    // A fixpoint, so it doesn't rely on the order `get_changes` returns.
    let mut excluded: HashSet<ChangeHash> = changes
        .iter()
        .filter(|(_, _, own_missing)| *own_missing)
        .map(|(hash, _, _)| *hash)
        .collect();
    let own_count = excluded.len();
    loop {
        let before = excluded.len();
        for (hash, deps, _) in &changes {
            if !excluded.contains(hash) && deps.iter().any(|dep| excluded.contains(dep)) {
                excluded.insert(*hash);
            }
        }
        if excluded.len() == before {
            break;
        }
    }

    // Heads of the kept part: kept changes no kept change depends on.
    let kept: Vec<&(ChangeHash, Vec<ChangeHash>, bool)> = changes
        .iter()
        .filter(|(hash, _, _)| !excluded.contains(hash))
        .collect();
    let depended_on: HashSet<ChangeHash> = kept
        .iter()
        .flat_map(|(_, deps, _)| deps.iter().copied())
        .collect();
    let mut before: Vec<ChangeHash> = kept
        .iter()
        .map(|(hash, _, _)| *hash)
        .filter(|hash| !depended_on.contains(hash))
        .collect();
    before.sort();

    Missing {
        before,
        own: own_count,
    }
}

/// The minimal set of store paths covering every change between `before` and `after`:
/// sorted, and no path is a prefix of another.
pub(crate) fn changed_paths(
    doc: &mut AutoCommit,
    before: &[ChangeHash],
    after: &[ChangeHash],
) -> Vec<Vec<String>> {
    let mut paths: Vec<Vec<String>> = doc.diff(before, after).iter().map(patch_path).collect();
    paths.sort();
    paths.dedup();

    // Sorted order puts a prefix right before its extensions, so comparing with the
    // last kept path is enough.
    let mut minimal: Vec<Vec<String>> = Vec::new();
    for path in paths {
        if minimal.last().is_some_and(|kept| path.starts_with(kept)) {
            continue;
        }
        minimal.push(path);
    }
    minimal
}

fn patch_path(patch: &automerge::Patch) -> Vec<String> {
    let mut path = Vec::new();
    for (_, prop) in &patch.path {
        match prop {
            Prop::Map(key) => path.push(key.clone()),
            // Inside a list: the list is replayed as a whole.
            Prop::Seq(_) => return path,
        }
    }

    match &patch.action {
        PatchAction::PutMap { key, .. } | PatchAction::DeleteMap { key, .. } => {
            path.push(key.clone())
        }
        PatchAction::Increment {
            prop: Prop::Map(key),
            ..
        } => path.push(key.clone()),
        // List and text edits: the containing object.
        _ => {}
    }
    path
}

#[cfg(test)]
mod tests {
    use super::*;
    use automerge::{transaction::Transactable, ObjType, ReadDoc, ScalarValue, ROOT};
    use serde_json::json;

    fn view(doc: &AutoCommit) -> serde_json::Value {
        store::read_path(doc, &[]).unwrap().unwrap()
    }

    #[test]
    fn minimal_paths() {
        let mut doc = AutoCommit::new();
        doc.put(ROOT, "untouched", 1_i64).unwrap();
        let before = doc.get_heads();

        let m = doc.put_object(ROOT, "m", ObjType::Map).unwrap();
        doc.put(&m, "a", 1_i64).unwrap();
        doc.put(ROOT, "x", 1_i64).unwrap();
        doc.put(ROOT, "c", ScalarValue::counter(0)).unwrap();
        doc.commit();
        let middle = doc.get_heads();
        doc.increment(ROOT, "c", 2).unwrap();
        let list = doc.put_object(ROOT, "l", ObjType::List).unwrap();
        doc.insert(&list, 0, "e").unwrap();
        doc.commit();
        let after = doc.get_heads();

        let all = changed_paths(&mut doc, &before, &after);
        assert_eq!(
            all,
            vec![
                vec!["c".to_owned()],
                vec!["l".to_owned()],
                vec!["m".to_owned()],
                vec!["x".to_owned()],
            ]
        );

        let later = changed_paths(&mut doc, &middle, &after);
        assert_eq!(later, vec![vec!["c".to_owned()], vec!["l".to_owned()]]);
    }

    #[test]
    fn replays_state_per_key() {
        // Base shared by the old document and the snapshot.
        let mut old = AutoCommit::new();
        let m = old.put_object(ROOT, "m", ObjType::Map).unwrap();
        old.put(&m, "keep", "base").unwrap();
        old.put(&m, "change", "base").unwrap();
        old.put(ROOT, "gone", true).unwrap();
        old.commit();
        let base = old.get_heads();

        let mut new = AutoCommit::new();
        store::materialize_into(&old, &mut new, &|_| false).unwrap();
        // Something the new epoch already holds that the delta must not clobber.
        let Some((_, m_new)) = new.get(ROOT, "m").unwrap() else {
            panic!("m missing");
        };
        new.put(&m_new, "keep", "new epoch").unwrap();
        new.commit();

        // The delta: one key changed, one key deleted.
        old.put(&m, "change", "delta").unwrap();
        old.delete(ROOT, "gone").unwrap();
        old.commit();
        let after = old.get_heads();

        let stats = replay(&mut old, &base, &after, &mut new);
        assert_eq!(stats.failed, 0);
        assert_eq!(
            view(&new),
            json!({ "m": { "keep": "new epoch", "change": "delta" } })
        );
    }

    #[test]
    fn missing_counts_only_own_changes() {
        let own = ActorId::from(&b"own-actor"[..]);
        let other = ActorId::from(&b"other-actor"[..]);

        let mut doc = AutoCommit::new().with_actor(own.clone());
        doc.put(ROOT, "base", 1_i64).unwrap();
        doc.commit();
        let base = doc.get_heads();

        // The snapshot contains the base only.
        let clock: BTreeMap<String, u64> = [(own.to_string(), 1)].into_iter().collect();

        // A foreign change after the base (replayed by its own node) ...
        let mut foreign = doc.fork().with_actor(other);
        foreign.put(ROOT, "theirs", 1_i64).unwrap();
        foreign.commit();
        doc.merge(&mut foreign).unwrap();
        let with_foreign = doc.get_heads();

        // ... and one of ours.
        doc.put(ROOT, "mine", 1_i64).unwrap();
        doc.commit();
        let after = doc.get_heads();

        let missing = missing_from_snapshot(&mut doc, &clock, &own);
        assert_eq!(missing.own, 1);
        assert_eq!(missing.before, with_foreign);
        assert_ne!(missing.before, base);
        assert_eq!(
            changed_paths(&mut doc, &missing.before, &after),
            vec![vec!["mine".to_owned()]]
        );

        // Nothing of ours missing: nothing to replay.
        let full: BTreeMap<String, u64> = [(own.to_string(), 2)].into_iter().collect();
        let none = missing_from_snapshot(&mut doc, &full, &own);
        assert_eq!(none.own, 0);
        let mut heads = after.clone();
        heads.sort();
        assert_eq!(none.before, heads);
    }

    #[test]
    fn protocol_metadata_is_never_replayed() {
        let mut old = AutoCommit::new();
        old.put(ROOT, "x", 1_i64).unwrap();
        old.commit();
        let base = old.get_heads();

        let mut new = AutoCommit::new();
        store::write_path(
            &mut new,
            &[META_ROOT.into(), EPOCHS.into(), "1".into()],
            &json!({ "generation": 1 }),
        )
        .unwrap();
        new.commit();

        // The old epoch replaces `_continuum` wholesale and writes a compaction record.
        store::write_path(
            &mut old,
            &[META_ROOT.into()],
            &json!({ "forgotten": { "pk": 1 } }),
        )
        .unwrap();
        store::write_path(
            &mut old,
            &[META_ROOT.into(), COMPACTIONS.into(), "c".into()],
            &json!({ "x": 1 }),
        )
        .unwrap();
        old.commit();
        let after = old.get_heads();

        replay(&mut old, &base, &after, &mut new);
        let meta = &view(&new)[META_ROOT];
        assert_eq!(meta[EPOCHS]["1"]["generation"], json!(1));
        assert_eq!(meta["forgotten"]["pk"], json!(1));
        assert!(meta.get(COMPACTIONS).is_none());
    }
}