//! Deterministic snapshot of the document at agreed base heads.
//!
//! Every node builds the snapshot of a committed compaction on its own. For that to
//! work, the snapshot must be byte-identical everywhere: same operations, same order,
//! same actor, same timestamp. Its change hash then identifies the new epoch, and a
//! node whose build produces a different hash knows it must fetch the snapshot instead.
//!
//! Determinism rests on:
//! * the content: the merged view at the base heads (`store::materialize_into`), which
//!   depends only on the base, not on how a node's history got there;
//! * a fixed actor derived from the compaction ID and generation;
//! * commit time 0 and no message;
//! * canonical (sorted) JSON for the metadata;
//! * the same Automerge version on every node (encodings may differ across versions).
//!
//! The snapshot carries its own metadata under `_continuum/epochs/<generation>`:
//! previous epoch, base heads, per-actor clock and retention policy. Later steps read it
//! at the snapshot change itself, so edits made afterwards through the store API can't
//! change what the protocol sees.

// `Snapshot::clock` is kept for returning nodes (step 7); late nodes read the clock from
// the metadata instead.
#![cfg_attr(not(test), allow(dead_code))]

use std::collections::BTreeMap;

use automerge::{transaction::CommitOptions, ActorId, AutoCommit, ChangeHash};
use serde_json::{json, Map as JsonMap, Value as JsonValue};
use sha2::{Digest, Sha256};

use super::records::{EpochRecord, Policy};
use crate::replica::Epoch;
use crate::store::{self, StoreError};

pub(crate) const META_ROOT: &str = "_continuum";
pub(crate) const COMPACTIONS: &str = "compactions";
pub(crate) const EPOCHS: &str = "epochs";

/// Everything that determines a snapshot besides the base content.
pub(crate) struct SnapshotParams<'a> {
    /// ID of the committed compaction; seeds the snapshot actor.
    pub(crate) compaction_id: &'a str,
    /// The epoch being compacted. The snapshot starts generation `previous + 1`.
    pub(crate) previous: Epoch,
    /// Heads of the base. Every node must hold all of these changes.
    pub(crate) base_heads: &'a [ChangeHash],
    /// Retention policy from the compaction request; stored verbatim (canonicalized).
    pub(crate) policy: &'a JsonValue,
    /// When the compaction was proposed (unix ms); recorded for replay windows.
    pub(crate) proposed_at: u64,
}

pub(crate) struct Snapshot {
    /// The new epoch: next generation, ID = hash of the snapshot change.
    pub(crate) epoch: Epoch,
    /// A document consisting of exactly that one change.
    pub(crate) doc: AutoCommit,
    /// Highest sequence number per actor (hex) contained in the base. A change
    /// (actor, seq) is part of the snapshot iff seq <= clock[actor].
    pub(crate) clock: BTreeMap<String, u64>,
}

#[derive(Debug, thiserror::Error)]
pub(crate) enum SnapshotError {
    #[error("base head {0} is not in the local document")]
    MissingBase(ChangeHash),
    #[error("snapshot produced {0} heads instead of exactly one")]
    UnexpectedHeads(usize),
    #[error(transparent)]
    Doc(#[from] automerge::AutomergeError),
    #[error(transparent)]
    Store(#[from] StoreError),
    #[error("snapshot task failed: {0}")]
    Task(String),
}

/// Owned version of [`SnapshotParams`], for moving into a blocking task.
#[derive(Clone, Debug)]
pub(crate) struct SnapshotRequest {
    pub(crate) compaction_id: String,
    pub(crate) previous: Epoch,
    pub(crate) base_heads: Vec<ChangeHash>,
    pub(crate) policy: JsonValue,
    pub(crate) proposed_at: u64,
}

impl SnapshotRequest {
    pub(crate) fn params(&self) -> SnapshotParams<'_> {
        SnapshotParams {
            compaction_id: &self.compaction_id,
            previous: self.previous,
            base_heads: &self.base_heads,
            policy: &self.policy,
            proposed_at: self.proposed_at,
        }
    }
}

/// Builds the snapshot from `copy`, a clone of the live document, on the blocking pool.
/// No lock is held while building; the caller only needs a lock for the clone.
/// Returns the copy as well, for callers that continue working with it.
pub(crate) async fn build_detached(
    mut copy: AutoCommit,
    request: SnapshotRequest,
) -> Result<(Snapshot, AutoCommit), SnapshotError> {
    tokio::task::spawn_blocking(move || {
        let snapshot = build(&mut copy, &request.params())?;
        Ok((snapshot, copy))
    })
    .await
    .map_err(|error| SnapshotError::Task(error.to_string()))?
}

/// The protocol's view of an epoch, as recorded in its snapshot.
#[derive(Clone, Debug)]
pub(crate) struct EpochMetadata {
    pub(crate) compaction_id: String,
    pub(crate) previous: Epoch,
    /// Highest sequence number per actor of `previous` contained in the snapshot.
    pub(crate) clock: BTreeMap<String, u64>,
    pub(crate) policy: Policy,
    /// When the compaction was proposed (unix ms): the point the cluster moved on.
    pub(crate) proposed_at: u64,
}

#[derive(serde::Deserialize)]
#[serde(rename_all = "camelCase")]
struct MetadataRecord {
    generation: u64,
    compaction_id: String,
    previous: EpochRecord,
    clock: BTreeMap<String, u64>,
    policy: Policy,
    #[serde(default)]
    proposed_at: u64,
}

fn parse_metadata(doc: &AutoCommit, generation: u64) -> Result<EpochMetadata, String> {
    let value = store::read_path(doc, &epoch_metadata_path(generation))
        .map_err(|error| error.to_string())?
        .ok_or_else(|| format!("no epoch metadata for generation {generation}"))?;
    let record: MetadataRecord = serde_json::from_value(value)
        .map_err(|error| format!("invalid epoch metadata: {error}"))?;
    if record.generation != generation {
        return Err(format!(
            "epoch metadata names generation {} instead of {generation}",
            record.generation
        ));
    }
    let previous = record
        .previous
        .to_epoch()
        .ok_or_else(|| "invalid previous epoch in metadata".to_owned())?;

    Ok(EpochMetadata {
        compaction_id: record.compaction_id,
        previous,
        clock: record.clock,
        policy: record.policy,
        proposed_at: record.proposed_at,
    })
}

/// Checks that `doc` is a document of `epoch`: its history starts with exactly the
/// snapshot change `epoch.id` (no dependencies), and the metadata at that change names
/// the same generation. Returns the metadata.
///
/// This does not prove the snapshot's content matches the old epoch; the hash commits
/// to it, and the peer is a cluster member trusted as much as for ordinary sync.
pub(crate) fn verify_document(doc: &mut AutoCommit, epoch: Epoch) -> Result<EpochMetadata, String> {
    let root = ChangeHash(epoch.id);
    let Some(change) = doc.get_change_by_hash(&root) else {
        return Err(format!(
            "document does not contain the snapshot change of {epoch}"
        ));
    };
    if !change.deps().is_empty() {
        return Err(format!("snapshot change of {epoch} has dependencies"));
    }

    let at_snapshot = doc.fork_at(&[root]).map_err(|error| error.to_string())?;
    parse_metadata(&at_snapshot, epoch.generation)
}

/// Metadata of an earlier generation, as the current document holds it. Snapshots
/// carry the metadata of all earlier epochs forward (only compaction records stay
/// behind), so a node many epochs behind can still find the epoch that followed its own.
pub(crate) fn read_epoch_metadata(
    doc: &AutoCommit,
    generation: u64,
) -> Result<EpochMetadata, String> {
    parse_metadata(doc, generation)
}

/// Path of the metadata record for an epoch generation.
pub(crate) fn epoch_metadata_path(generation: u64) -> Vec<String> {
    vec![
        META_ROOT.to_owned(),
        EPOCHS.to_owned(),
        generation.to_string(),
    ]
}

/// Builds the snapshot of `source` at `params.base_heads`.
pub(crate) fn build(
    source: &mut AutoCommit,
    params: &SnapshotParams<'_>,
) -> Result<Snapshot, SnapshotError> {
    // fork_at with unknown heads would quietly build something else.
    for head in params.base_heads {
        if source.get_change_by_hash(head).is_none() {
            return Err(SnapshotError::MissingBase(*head));
        }
    }

    let mut base = source.fork_at(params.base_heads)?;
    let clock = clock_of(&mut base);
    let generation = params.previous.generation + 1;

    let mut snapshot =
        AutoCommit::new().with_actor(snapshot_actor(params.compaction_id, generation));

    // Compaction records describe the old epoch's protocol and stay behind.
    store::materialize_into(&base, &mut snapshot, &is_compaction_record)?;

    let mut base_heads: Vec<String> = params.base_heads.iter().map(|h| h.to_string()).collect();
    base_heads.sort();

    let metadata = canonical(json!({
        "generation": generation,
        "compactionId": params.compaction_id,
        "previous": {
            "generation": params.previous.generation,
            "id": params.previous.id_hex(),
        },
        "baseHeads": base_heads,
        "clock": clock,
        "policy": params.policy,
        "proposedAt": params.proposed_at,
    }));
    store::write_path(&mut snapshot, &epoch_metadata_path(generation), &metadata)?;

    let _ = snapshot.commit_with(CommitOptions::default().with_time(0));

    let heads = snapshot.get_heads();
    let [id] = heads.as_slice() else {
        return Err(SnapshotError::UnexpectedHeads(heads.len()));
    };

    Ok(Snapshot {
        epoch: Epoch {
            generation,
            id: id.0,
        },
        doc: snapshot,
        clock,
    })
}

/// Highest sequence number per actor in `doc`'s history.
fn clock_of(doc: &mut AutoCommit) -> BTreeMap<String, u64> {
    let mut clock = BTreeMap::new();
    for change in doc.get_changes(&[]) {
        let seq = clock.entry(change.actor_id().to_string()).or_insert(0);
        *seq = (*seq).max(change.seq());
    }
    clock
}

/// The same 16-byte actor on every node, unique per compaction and generation.
fn snapshot_actor(compaction_id: &str, generation: u64) -> ActorId {
    let digest = Sha256::new()
        .chain_update(b"continuum.snapshot\0")
        .chain_update(compaction_id.as_bytes())
        .chain_update(b"\0")
        .chain_update(generation.to_be_bytes())
        .finalize();
    ActorId::from(&digest[..16])
}

fn is_compaction_record(path: &[String]) -> bool {
    matches!(path, [root, records] if root == META_ROOT && records == COMPACTIONS)
}

/// Rebuilds every object with sorted keys, so the result is the same whichever key
/// order serde_json happens to use (it keeps insertion order with `preserve_order`).
fn canonical(value: JsonValue) -> JsonValue {
    match value {
        JsonValue::Object(map) => {
            let sorted: BTreeMap<String, JsonValue> = map
                .into_iter()
                .map(|(key, value)| (key, canonical(value)))
                .collect();
            JsonValue::Object(sorted.into_iter().collect::<JsonMap<_, _>>())
        }
        JsonValue::Array(items) => JsonValue::Array(items.into_iter().map(canonical).collect()),
        other => other,
    }
}

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

    fn doc(actor: &[u8]) -> AutoCommit {
        AutoCommit::new().with_actor(ActorId::from(actor))
    }

    fn params<'a>(heads: &'a [ChangeHash], policy: &'a JsonValue) -> SnapshotParams<'a> {
        SnapshotParams {
            compaction_id: "c-1",
            previous: Epoch::GENESIS,
            base_heads: heads,
            policy,
            proposed_at: 0,
        }
    }

    fn snapshot_of(source: &mut AutoCommit) -> Snapshot {
        let heads = source.get_heads();
        let policy = json!({});
        build(source, &params(&heads, &policy)).expect("snapshot")
    }

    /// The user-visible view, without our own metadata.
    fn user_view(doc: &AutoCommit) -> JsonValue {
        let mut view = store::read_path(doc, &[]).unwrap().unwrap();
        view.as_object_mut().unwrap().remove(META_ROOT);
        view
    }

    #[test]
    fn same_view_different_history_same_snapshot() {
        // One history: overwrites, increments, one-shot text.
        let mut a = doc(b"actor-a");
        a.put(ROOT, "x", 1_i64).unwrap();
        a.commit();
        a.put(ROOT, "x", 2_i64).unwrap();
        a.put(ROOT, "c", ScalarValue::counter(0)).unwrap();
        a.increment(ROOT, "c", 5).unwrap();
        let text = a.put_object(ROOT, "t", ObjType::Text).unwrap();
        a.splice_text(&text, 0, 0, "hello").unwrap();
        let m = a.put_object(ROOT, "m", ObjType::Map).unwrap();
        a.put(&m, "k", "v").unwrap();
        a.commit();

        // Another actor, other order, other steps, same result.
        let mut b = doc(b"actor-b");
        let m = b.put_object(ROOT, "m", ObjType::Map).unwrap();
        b.put(&m, "k", "v").unwrap();
        let text = b.put_object(ROOT, "t", ObjType::Text).unwrap();
        b.splice_text(&text, 0, 0, "hel").unwrap();
        b.commit();
        b.splice_text(&text, 3, 0, "lo").unwrap();
        b.put(ROOT, "c", ScalarValue::counter(5)).unwrap();
        b.put(ROOT, "x", 2_i64).unwrap();
        b.commit();

        assert_eq!(user_view(&a), user_view(&b));

        let snap_a = snapshot_of(&mut a);
        let snap_b = snapshot_of(&mut b);

        // Different histories mean different clocks, and the clock is part of the
        // metadata. Equal views alone therefore give equal content but not equal IDs;
        // equal bases do give equal IDs (see `rebuild_is_deterministic`).
        assert_eq!(user_view(&snap_a.doc), user_view(&snap_b.doc));
        assert_eq!(user_view(&snap_a.doc), user_view(&a));
    }

    #[test]
    fn rebuild_is_deterministic() {
        let mut a = doc(b"actor-a");
        a.put(ROOT, "x", 1_i64).unwrap();
        let list = a.put_object(ROOT, "l", ObjType::List).unwrap();
        a.insert(&list, 0, "first").unwrap();
        let inner = a.insert_object(&list, 1, ObjType::Map).unwrap();
        a.put(&inner, "k", true).unwrap();
        a.commit();

        // A second node holding the same changes, received in a different way.
        let mut b = doc(b"actor-b");
        b.merge(&mut a.clone()).unwrap();

        let first = snapshot_of(&mut a);
        let again = snapshot_of(&mut a);
        let elsewhere = snapshot_of(&mut b);

        assert_eq!(first.epoch, again.epoch);
        assert_eq!(first.epoch, elsewhere.epoch);
        assert_eq!(first.epoch.generation, 1);
    }

    #[test]
    fn conflicts_collapse_into_the_merged_view() {
        let mut a = doc(b"actor-a");
        let cfg = a.put_object(ROOT, "cfg", ObjType::Map).unwrap();
        a.put(&cfg, "x", 1_i64).unwrap();
        a.commit();

        let mut b = doc(b"actor-b");
        let cfg = b.put_object(ROOT, "cfg", ObjType::Map).unwrap();
        b.put(&cfg, "y", 2_i64).unwrap();
        b.commit();

        a.merge(&mut b).unwrap();
        let (_, conflicts) = store::read_path_with_conflicts(&a, &[]).unwrap().unwrap();
        assert!(
            !conflicts.is_empty(),
            "test setup should produce a conflict"
        );

        let snap = snapshot_of(&mut a);
        assert_eq!(user_view(&snap.doc), json!({ "cfg": { "x": 1, "y": 2 } }));

        let (_, conflicts) = store::read_path_with_conflicts(&snap.doc, &[])
            .unwrap()
            .unwrap();
        assert!(
            conflicts.is_empty(),
            "snapshot still has conflicts: {conflicts:?}"
        );
    }

    #[test]
    fn base_heads_bound_content_and_clock() {
        let mut a = doc(b"actor-a");
        for i in 0..3_i64 {
            a.put(ROOT, format!("a{i}"), i).unwrap();
            a.commit();
        }

        let mut b = doc(b"actor-b");
        for i in 0..2_i64 {
            b.put(ROOT, format!("b{i}"), i).unwrap();
            b.commit();
        }
        a.merge(&mut b).unwrap();

        let base = a.get_heads();

        // After the base: must not appear in the snapshot or its clock.
        a.put(ROOT, "late", true).unwrap();
        a.commit();

        let policy = json!({});
        let snap = build(&mut a, &params(&base, &policy)).unwrap();

        let view = user_view(&snap.doc);
        assert!(view.get("late").is_none());
        assert_eq!(view["a2"], json!(2));
        assert_eq!(view["b1"], json!(1));

        let actor_a = ActorId::from(b"actor-a".as_slice()).to_string();
        let actor_b = ActorId::from(b"actor-b".as_slice()).to_string();
        assert_eq!(snap.clock.get(&actor_a), Some(&3));
        assert_eq!(snap.clock.get(&actor_b), Some(&2));
    }

    #[test]
    fn metadata_written_and_compaction_records_dropped() {
        let mut a = doc(b"actor-a");
        store::write_path(
            &mut a,
            &[META_ROOT.into(), COMPACTIONS.into(), "c-1".into()],
            &json!({ "state": "committed" }),
        )
        .unwrap();
        store::write_path(
            &mut a,
            &[META_ROOT.into(), "forgotten".into(), "pk".into()],
            &json!({ "at": 1 }),
        )
        .unwrap();
        a.commit();

        let heads = a.get_heads();
        let policy = json!({ "replayWindow": { "epochs": 1 } });
        let snap = build(&mut a, &params(&heads, &policy)).unwrap();

        let meta = store::read_path(&snap.doc, &[META_ROOT.into()])
            .unwrap()
            .unwrap();
        assert!(meta.get(COMPACTIONS).is_none());
        assert_eq!(meta["forgotten"]["pk"]["at"], json!(1));

        let epoch = &meta[EPOCHS]["1"];
        assert_eq!(epoch["generation"], json!(1));
        assert_eq!(epoch["compactionId"], json!("c-1"));
        assert_eq!(epoch["previous"]["generation"], json!(0));
        assert_eq!(epoch["policy"], policy);
        assert_eq!(epoch["baseHeads"].as_array().unwrap().len(), heads.len());
    }

    #[test]
    fn automerge_types_survive() {
        let mut a = doc(b"actor-a");
        a.put(ROOT, "c", ScalarValue::counter(3)).unwrap();
        a.increment(ROOT, "c", 4).unwrap();
        let text = a.put_object(ROOT, "t", ObjType::Text).unwrap();
        a.splice_text(&text, 0, 0, "abc").unwrap();
        a.put(ROOT, "bytes", ScalarValue::Bytes(vec![0, 1, 255]))
            .unwrap();
        a.put(ROOT, "ts", ScalarValue::Timestamp(1_700_000_000_000))
            .unwrap();
        a.put(ROOT, "neg_zero", -0.0_f64).unwrap();
        a.commit();

        let snap = snapshot_of(&mut a);
        let d = &snap.doc;

        match d.get(ROOT, "c").unwrap() {
            Some((Value::Scalar(s), _)) => match s.as_ref() {
                ScalarValue::Counter(c) => assert_eq!(i64::from(c.clone()), 7),
                other => panic!("counter became {other:?}"),
            },
            other => panic!("unexpected {other:?}"),
        }

        let (_, text_id) = d.get(ROOT, "t").unwrap().unwrap();
        assert_eq!(d.object_type(&text_id).unwrap(), ObjType::Text);
        assert_eq!(d.text(&text_id).unwrap(), "abc");

        match d.get(ROOT, "bytes").unwrap() {
            Some((Value::Scalar(s), _)) => {
                assert_eq!(s.as_ref(), &ScalarValue::Bytes(vec![0, 1, 255]))
            }
            other => panic!("unexpected {other:?}"),
        }
        match d.get(ROOT, "ts").unwrap() {
            Some((Value::Scalar(s), _)) => {
                assert_eq!(s.as_ref(), &ScalarValue::Timestamp(1_700_000_000_000))
            }
            other => panic!("unexpected {other:?}"),
        }
        match d.get(ROOT, "neg_zero").unwrap() {
            Some((Value::Scalar(s), _)) => match s.as_ref() {
                ScalarValue::F64(f) => assert_eq!(f.to_bits(), (-0.0_f64).to_bits()),
                other => panic!("float became {other:?}"),
            },
            other => panic!("unexpected {other:?}"),
        }
    }

    #[test]
    fn snapshot_is_a_single_change_with_the_epoch_id() {
        let mut a = doc(b"actor-a");
        a.put(ROOT, "x", 1_i64).unwrap();
        a.commit();

        let mut snap = snapshot_of(&mut a);
        let heads = snap.doc.get_heads();
        assert_eq!(heads.len(), 1);
        assert_eq!(heads[0].0, snap.epoch.id);
        assert_eq!(snap.doc.get_changes(&[]).len(), 1);
    }

    #[test]
    fn missing_base_is_rejected() {
        let mut a = doc(b"actor-a");
        a.put(ROOT, "x", 1_i64).unwrap();
        a.commit();

        let mut other = doc(b"actor-b");
        other.put(ROOT, "y", 1_i64).unwrap();
        other.commit();
        let foreign = other.get_heads();

        let policy = json!({});
        assert!(matches!(
            build(&mut a, &params(&foreign, &policy)),
            Err(SnapshotError::MissingBase(_))
        ));
    }

    #[test]
    fn empty_document() {
        let mut a = doc(b"actor-a");
        let snap = snapshot_of(&mut a);
        assert_eq!(snap.epoch.generation, 1);
        assert!(snap.clock.is_empty());
        assert_eq!(user_view(&snap.doc), json!({}));
    }

    #[test]
    fn fetched_document_is_verified_and_late_changes_replayed() {
        use crate::compaction::delta;

        // Old epoch: the base, then a late write by this node.
        let mut late = doc(b"late-node");
        late.put(ROOT, "shared", "base").unwrap();
        late.commit();
        let base = late.get_heads();
        let policy = json!({});
        let mut snap = build(&mut late, &params(&base, &policy)).unwrap();

        // Meanwhile the cluster moved on in the new epoch.
        let mut cluster = snap
            .doc
            .fork()
            .with_actor(ActorId::from(&b"cluster-node"[..]));
        cluster.put(ROOT, "new", 1_i64).unwrap();
        cluster.commit();
        let bytes = cluster.save();

        late.put(ROOT, "mine", "late").unwrap();
        late.commit();

        let mut fetched = AutoCommit::load(&bytes).unwrap();
        let metadata = verify_document(&mut fetched, snap.epoch).unwrap();
        assert_eq!(metadata.previous, Epoch::GENESIS);
        assert_eq!(metadata.compaction_id, "c-1");
        assert_eq!(
            metadata.policy.replay_window,
            Policy::default_replay_window()
        );

        // Documents of other epochs are refused.
        let mut wrong = snap.epoch;
        wrong.generation = 2;
        assert!(verify_document(&mut fetched, wrong).is_err());
        let mut wrong = snap.epoch;
        wrong.id = [0; 32];
        assert!(verify_document(&mut fetched, wrong).is_err());

        let own = late.get_actor().clone();
        let missing = delta::missing_from_snapshot(&mut late, &metadata.clock, &own);
        assert_eq!(missing.own, 1);
        let heads = late.get_heads();
        let mut adopted = fetched.with_actor(ActorId::from(&b"late-node-next"[..]));
        delta::replay(&mut late, &missing.before, &heads, &mut adopted);
        assert_eq!(
            user_view(&adopted),
            json!({ "shared": "base", "new": 1, "mine": "late" })
        );
    }
}