//! One Automerge sync with a peer, from a reset to the next.
//!
//! Automerge's sync state (`automerge::sync::State`) is private to this module: the
//! rest of the node generates, receives and inspects through `SyncSession`'s methods.
//! The workarounds for its assumption of reliable delivery (resending after a timeout,
//! forgetting what was sent when the peer still needs it) live here with it, so an
//! Automerge upgrade that changes the state touches this file only.

use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;

use automerge::sync::{self, SyncDoc};
use automerge::{AutoCommit, AutomergeError};
use tokio::time::Instant;
use tracing::debug;

use crate::backoff::Backoff;
use crate::swarm::throughput::SyncBytes;

/// A sync message being handed to our transport (for Mycelium: the HTTP request is in
/// progress).
#[derive(Clone)]
pub(crate) struct Sending {
    /// Changes it carries for the first time; taken back if the send fails.
    pub(crate) sent: Vec<automerge::ChangeHash>,
    /// A sync pass was requested meanwhile. Once this one is handed over, the peer must
    /// be enqueued again so the deferred changes go out.
    pub(crate) regenerate: bool,
}

/// Session IDs, unique within the process.
static NEXT_SESSION: AtomicU64 = AtomicU64::new(1);

/// First wait before generating again after our transport refused a message; doubles
/// while it keeps refusing.
const SEND_RETRY: Duration = Duration::from_millis(500);
const SEND_RETRY_MAX: Duration = Duration::from_secs(30);

/// One Automerge sync with a peer, from a reset to the next. Everything a reset must
/// forget lives here, so a reset replaces it whole and can't miss a field.
#[derive(Clone)]
pub(crate) struct SyncSession {
    /// Tells this session from the one a reset puts in its place (`NEXT_SESSION`).
    id: u64,
    state: sync::State,

    /// A message being handed to our transport.
    pub(crate) sending: Option<Sending>,

    /// After our transport refused a message: no new one before this. The refused one is
    /// not kept; the next one is generated afresh and carries everything pending.
    pub(crate) retry_at: Option<Instant>,
    /// The wait after the next refusal.
    pub(crate) send_backoff: Backoff,

    /// We offered the peer a delta instead of a message too large to send
    /// (`withdraw_for_delta`): no more messages until it answers, or the answer is
    /// overdue (then the offer goes out again).
    pub(crate) delta_offered: bool,
    /// What the last `generate` changed in the sync state besides the hashes sent: the
    /// heads it last sent, and whether it had answered yet. Restored when the message is
    /// withdrawn for a delta: otherwise Automerge believes the peer heard our heads, and
    /// once the peer has the delta and says its heads are ours, never answers it, so the
    /// peer never learns our heads (a joining node then never finishes its initial
    /// sync).
    before_generate: (Vec<automerge::ChangeHash>, bool),

    /// The answer we expect to our last sync message(s), if any.
    pub(crate) reply: Option<SyncReplyWait>,
    /// How long to wait for it; doubles while answers stay out.
    pub(crate) reply_backoff: Backoff,
}

/// Sync messages we sent and expect an answer to. Automerge assumes reliable delivery:
/// a lost message leaves both sides waiting (`in_flight`), and changes it carried are
/// never sent again (`sent_hashes`). We resend after a timeout instead.
#[derive(Clone, Debug)]
pub(crate) struct SyncReplyWait {
    pub(crate) deadline: Instant,
    /// When the first of the unanswered messages was generated, and the sync bytes
    /// exchanged with the peer by then (`swarm::throughput`).
    pub(crate) started: (Instant, SyncBytes),
    /// Changes first sent in the unanswered messages.
    pub(crate) sent: Vec<automerge::ChangeHash>,
}

impl SyncSession {
    /// `reply_backoff`: waits for sync answers (`Timing::sync_reply_backoff`).
    pub(crate) fn new(reply_backoff: Backoff) -> Self {
        Self {
            id: NEXT_SESSION.fetch_add(1, Ordering::Relaxed),
            state: sync::State::new(),
            sending: None,
            retry_at: None,
            send_backoff: Backoff::new(SEND_RETRY, SEND_RETRY_MAX),
            delta_offered: false,
            before_generate: (Vec::new(), false),
            reply: None,
            reply_backoff,
        }
    }

    /// Whether the peer said anything in this session yet (its heads are known).
    pub(crate) fn has_opened(&self) -> bool {
        self.state.their_heads.is_some()
    }

    /// The next sync message for the peer from `doc`, if there is anything to say. A
    /// message is marked as being handed over (`sending`), and an answer to it is
    /// expected, unless the peer's heads are ours already: it has everything, and
    /// Automerge answers only when it has something to say, so waiting for an answer
    /// would only end in a pointless resend. `bytes`: the sync bytes exchanged with the
    /// peer so far (`swarm::throughput`).
    pub(crate) fn generate(
        &mut self,
        doc: &mut AutoCommit,
        now: Instant,
        bytes: SyncBytes,
    ) -> Option<sync::Message> {
        let sent_before = self.state.sent_hashes.clone();
        self.before_generate = (
            self.state.last_sent_heads.clone(),
            self.state.have_responded,
        );
        let message = doc.sync().generate_sync_message(&mut self.state);
        let in_step = self.state.their_heads.as_ref() == Some(&doc.get_heads());
        let sent: Vec<_> = self
            .state
            .sent_hashes
            .difference(&sent_before)
            .copied()
            .collect();
        if in_step {
            self.reply_settled();
        } else if message.is_some() {
            self.expect_reply(sent.clone(), now, bytes);
        }
        if message.is_some() {
            // Marked together with the generation, so no second message is generated
            // while this one is handed over.
            self.sending = Some(Sending {
                sent,
                regenerate: false,
            });
        }
        message
    }

    /// Applies a sync message from the peer to `doc`. Doesn't settle the awaited
    /// answer: the caller measures the exchange first (`reply`).
    pub(crate) fn receive(
        &mut self,
        doc: &mut AutoCommit,
        message: sync::Message,
    ) -> Result<(), AutomergeError> {
        self.delta_offered = false;
        doc.sync().receive_sync_message(&mut self.state, message)
    }

    /// The message just generated is too large to send in one piece: take back what it
    /// carried, as if it had never been generated, and offer a delta instead. The
    /// message being handed over becomes the offer (it carries nothing); the answer is
    /// still awaited, and if it doesn't come, the offer is made again. Returns the
    /// peer's heads as far as we know them: where the delta starts.
    pub(crate) fn withdraw_for_delta(&mut self) -> Vec<automerge::ChangeHash> {
        if let Some(sending) = &mut self.sending {
            let sent = std::mem::take(&mut sending.sent);
            for hash in &sent {
                self.state.sent_hashes.remove(hash);
            }
            if let Some(wait) = &mut self.reply {
                wait.sent.retain(|hash| !sent.contains(hash));
            }
        }
        self.state.in_flight = false;
        (self.state.last_sent_heads, self.state.have_responded) = self.before_generate.clone();
        self.delta_offered = true;
        self.state.their_heads.clone().unwrap_or_default()
    }

    /// Whether `doc` holds everything the peer had when it last told us its heads: they
    /// are known, not empty (empty heads are "contained" in every document, so they
    /// prove nothing), and all in `doc`.
    pub(crate) fn have_all_theirs(&self, doc: &mut AutoCommit) -> bool {
        self.state.their_heads.as_ref().is_some_and(|heads| {
            !heads.is_empty()
                && heads
                    .iter()
                    .all(|head| doc.get_change_by_hash(head).is_some())
        })
    }

    /// The peer has heads we lack (its keepalive said so): ask for what they need,
    /// even though Automerge thinks a message is still in flight.
    pub(crate) fn pull_from(&mut self, heads: Vec<automerge::ChangeHash>) {
        self.state.their_heads = Some(heads);
        self.state.in_flight = false;
    }

    /// Automerge never resends a change it believes the peer has received
    /// (`sent_hashes`). If the message carrying it was lost while later ones arrive, the
    /// peer keeps asking for it (`need`, from the heads we advertise) and we keep not
    /// sending it: a livelock the reply timeout can't see (answers keep coming) and the
    /// keepalive heal can't either (sync never goes idle). So when the peer needs a
    /// change we think it has, what we think it has is wrong: forget all of it, and the
    /// next message carries everything its bloom filter lacks. (A reordered message can
    /// make that send a change twice; cheap.) Call after `receive`.
    pub(crate) fn resend_if_needed(&mut self, peer: &str) {
        let Some(need) = &self.state.their_need else {
            return;
        };
        if need
            .iter()
            .any(|hash| self.state.sent_hashes.contains(hash))
        {
            debug!(
                peer = %peer,
                sent = self.state.sent_hashes.len(),
                "Peer needs changes we had sent; forgetting what was sent"
            );
            self.state.sent_hashes.clear();
        }
        if answer_for_everything_since_shared(&mut self.state) {
            debug!(
                peer = %peer,
                need = self.state.their_need.as_ref().map_or(0, Vec::len),
                "Peer lacks the changes ours build on; sending everything since what we share"
            );
        }
    }

    /// For the simulator's dump: the peer's heads as Automerge knows them, and whether
    /// it thinks a message is in flight.
    #[cfg(test)]
    pub(crate) fn automerge_view(&self) -> (Option<&[automerge::ChangeHash]>, bool) {
        (self.state.their_heads.as_deref(), self.state.in_flight)
    }

    /// Whether we leave the first message of this session to the peer.
    ///
    /// After a reset, both sides usually open at once, and the openings cross. Between
    /// peers that share history that costs nothing. But a peer with an empty document
    /// answers our opening with empty heads, which Automerge takes for "lost
    /// everything", so the whole document would go out a second time. So a ready node
    /// never opens towards a peer that isn't ready (one that is still catching up); that
    /// peer opens itself.
    pub(crate) fn leave_opening(&self, we_ready: bool, peer_ready: Option<bool>) -> bool {
        self.state.their_heads.is_none() && we_ready && peer_ready == Some(false)
    }

    /// Which session this is: a message generated from it reports its send's outcome
    /// with this (`send_succeeded`, `send_failed`).
    pub(crate) fn id(&self) -> u64 {
        self.id
    }

    /// Our transport took the message generated from session `session`. Returns whether
    /// to generate again: a pass was requested while it was handed over, or a reset
    /// replaced the session meanwhile (any pass requested then was lost with the old
    /// session, and the new one hasn't necessarily opened yet).
    pub(crate) fn send_succeeded(&mut self, session: u64) -> bool {
        if session != self.id {
            return true;
        }
        self.retry_at = None;
        self.send_backoff.reset();
        self.sending
            .take()
            .is_some_and(|sending| sending.regenerate)
    }

    /// Our transport refused the message: as far as the peer is concerned, it was never
    /// sent. Takes back its changes (they go out with the next message), clears
    /// `in_flight`, and returns how long to wait before generating again. None if a reset
    /// already discarded what was being sent (the message came from another `session`).
    pub(crate) fn send_failed(&mut self, session: u64, now: Instant) -> Option<Duration> {
        if session != self.id {
            return None;
        }
        let sending = self.sending.take()?;
        for hash in &sending.sent {
            self.state.sent_hashes.remove(hash);
        }
        if let Some(wait) = &mut self.reply {
            wait.sent.retain(|hash| !sending.sent.contains(hash));
        }
        self.state.in_flight = false;
        let delay = self.send_backoff.next();
        self.retry_at = Some(now + delay);
        Some(delay)
    }

    /// A sync message went out; `sent` are the changes it carried for the first time.
    /// `bytes` are the sync bytes exchanged with the peer before it.
    pub(crate) fn expect_reply(
        &mut self,
        sent: Vec<automerge::ChangeHash>,
        now: Instant,
        bytes: SyncBytes,
    ) {
        match &mut self.reply {
            // Still waiting for an earlier one: keep its deadline, remember these too.
            Some(wait) => wait.sent.extend(sent),
            None => {
                self.reply = Some(SyncReplyWait {
                    deadline: now + self.reply_backoff.current(),
                    started: (now, bytes),
                    sent,
                });
            }
        }
    }

    /// No answer is due any more: the peer answered, or keepalives show both documents
    /// are the same.
    pub(crate) fn reply_settled(&mut self) {
        self.reply = None;
        self.reply_backoff.reset();
    }

    /// If the answer is overdue: forgets that the unanswered changes were sent and that
    /// a message is in flight, so the next sync pass sends them again. Returns how many
    /// changes that concerns.
    pub(super) fn take_overdue_reply(&mut self, now: Instant) -> Option<usize> {
        if self.reply.as_ref().is_none_or(|wait| now < wait.deadline) {
            return None;
        }
        let wait = self.reply.take()?;
        for hash in &wait.sent {
            self.state.sent_hashes.remove(hash);
        }
        self.state.in_flight = false;
        self.delta_offered = false;
        self.reply_backoff.grow();
        Some(wait.sent.len())
    }
}

/// A peer that got a change of ours without the changes it builds on (they were lost,
/// a later one arrived) asks for those it lacks, and while it does, Automerge leaves
/// out what it has (its bloom filter: `have` waits until all it needs is in our heads).
/// Without a `have`, Automerge sends exactly what was asked for, and the peer, finding
/// each change's own predecessors missing in turn, asks again: one round trip per
/// missing change, hundreds after a short outage. So answer as if it had said it has
/// nothing since the heads we share: everything since then goes in one message (a
/// delta when large: `swarm::sync`), and what it had already only twice. Returns
/// whether it did. Call after receiving.
fn answer_for_everything_since_shared(state: &mut sync::State) -> bool {
    let asks_without_saying = state.their_have.as_ref().is_some_and(Vec::is_empty)
        && state
            .their_need
            .as_ref()
            .is_some_and(|need| !need.is_empty());
    if !asks_without_saying {
        return false;
    }
    state.their_have = Some(vec![sync::Have {
        last_sync: state.shared_heads.clone(),
        bloom: sync::BloomFilter::default(),
    }]);
    // What we sent since may be what was lost.
    state.sent_hashes.clear();
    true
}

#[cfg(test)]
mod tests {
    use super::*;

    /// After a loss, the changes the peer lacks go in one message, not one per round
    /// trip (`answer_for_everything_since_shared`).
    #[test]
    fn missing_predecessors_go_at_once() {
        use automerge::transaction::Transactable;
        let mut a = AutoCommit::new();
        let mut b = AutoCommit::new();
        let (mut at_a, mut at_b) = (sync::State::new(), sync::State::new());
        a.put(automerge::ROOT, "early", 1).unwrap();
        // Until in step.
        for _ in 0..10 {
            let to_b = a.sync().generate_sync_message(&mut at_a);
            let to_a = b.sync().generate_sync_message(&mut at_b);
            if to_b.is_none() && to_a.is_none() {
                break;
            }
            if let Some(m) = to_b {
                b.sync().receive_sync_message(&mut at_b, m).unwrap();
            }
            if let Some(m) = to_a {
                a.sync().receive_sync_message(&mut at_a, m).unwrap();
            }
        }
        assert_eq!(a.get_heads(), b.get_heads());
        // Five changes, whose message is lost; then a sixth, which arrives alone.
        for i in 0..5 {
            a.put(automerge::ROOT, format!("lost{i}"), i).unwrap();
            a.commit();
        }
        let _lost = a.sync().generate_sync_message(&mut at_a).unwrap();
        a.put(automerge::ROOT, "last", 6).unwrap();
        a.commit();
        at_a.in_flight = false; // As after our reply timeout.
        let last = a.sync().generate_sync_message(&mut at_a).unwrap();
        assert_eq!(last.changes.len(), 1);
        b.sync().receive_sync_message(&mut at_b, last).unwrap();
        // b asks for the fifth, saying nothing of what it has.
        let ask = b.sync().generate_sync_message(&mut at_b).unwrap();
        assert!(ask.have.is_empty() && ask.need.len() == 1);
        a.sync().receive_sync_message(&mut at_a, ask).unwrap();
        assert!(answer_for_everything_since_shared(&mut at_a));
        // All since the shared heads, in one message (here the whole document: more
        // than a third of it is missing).
        let answer = a.sync().generate_sync_message(&mut at_a).unwrap();
        b.sync().receive_sync_message(&mut at_b, answer).unwrap();
        assert_eq!(b.get_heads(), a.get_heads());
        // A message with what it has is left alone.
        assert!(!answer_for_everything_since_shared(&mut at_b));
    }

    fn session() -> SyncSession {
        SyncSession::new(Backoff::new(
            Duration::from_secs(1),
            Duration::from_secs(30),
        ))
    }

    fn sending() -> Option<Sending> {
        Some(Sending {
            sent: Vec::new(),
            regenerate: false,
        })
    }

    #[test]
    fn a_send_from_a_replaced_session_leaves_the_new_one_alone() {
        let old = session();
        let mut new = session();
        new.sending = sending();

        // A reset discarded the old session while its message was handed over: the
        // new one's own message stays, and the peer is synced again.
        assert!(new.send_succeeded(old.id()));
        assert!(new.sending.is_some());
        assert_eq!(new.send_failed(old.id(), Instant::now()), None);
        assert!(new.sending.is_some());
        assert_eq!(new.retry_at, None);

        // Its own outcome still applies.
        assert!(!new.send_succeeded(new.id()));
        assert!(new.sending.is_none());
    }

    /// A message too large to send is taken back: the peer, after the offer and an
    /// overdue answer, gets the same changes in the next message.
    #[test]
    fn withdrawing_for_a_delta_keeps_the_changes_for_later() {
        use automerge::{ROOT, transaction::Transactable};

        // 16 KiB of hex, about 8 KiB of information: Automerge compresses what it sends
        // (a peer with nothing gets the document saved), so this stays above 4 KB.
        let mut rng = crate::testrand::Rng::new(1);
        let big: String = (0..1024)
            .map(|_| format!("{:016x}", rng.next_u64()))
            .collect();
        let mut ours = AutoCommit::new();
        ours.put(ROOT, "big", big).unwrap();
        ours.commit();
        let mut theirs = AutoCommit::new();
        let (mut a, mut b) = (session(), session());
        let now = Instant::now();
        let bytes = SyncBytes::default();

        // a opens, b answers with its (empty) heads, then a sends the changes.
        let opening = a.generate(&mut ours, now, bytes).unwrap();
        a.send_succeeded(a.id());
        b.receive(&mut theirs, opening).unwrap();
        let answer = b.generate(&mut theirs, now, bytes).unwrap();
        b.send_succeeded(b.id());
        a.receive(&mut ours, answer).unwrap();
        let carrying = a.generate(&mut ours, now, bytes).unwrap();
        let carried = carrying.clone().encode().len();
        assert!(carried > 4_000, "the changes took {carried} bytes");

        let from = a.withdraw_for_delta();
        assert!(from.is_empty());
        assert!(a.delta_offered);
        assert!(!a.send_succeeded(a.id()));

        // No answer to the offer: it lapses, and the changes go out again.
        assert!(
            a.take_overdue_reply(now + Duration::from_secs(3600))
                .is_some()
        );
        assert!(!a.delta_offered);
        let again = a.generate(&mut ours, now, bytes).unwrap();
        let carried_again = again.encode().len();
        assert!(carried_again > 4_000, "sent again in {carried_again} bytes");
    }

    /// A joining node opens, the answer is too large and becomes a delta, the node gets
    /// the delta and says so: the answer to that carries our heads, so the node knows it
    /// has everything (`have_all_theirs`).
    #[test]
    fn after_a_delta_the_peer_learns_our_heads() {
        use automerge::{ROOT, transaction::Transactable};

        let mut rng = crate::testrand::Rng::new(2);
        let big: String = (0..1024)
            .map(|_| format!("{:016x}", rng.next_u64()))
            .collect();
        let mut ours = AutoCommit::new();
        ours.put(ROOT, "big", big).unwrap();
        ours.commit();
        let mut theirs = AutoCommit::new();
        let (mut a, mut b) = (session(), session());
        let now = Instant::now();
        let bytes = SyncBytes::default();

        // b (joining) opens; a's answer would carry everything: withdrawn for a delta.
        let opening = b.generate(&mut theirs, now, bytes).unwrap();
        b.send_succeeded(b.id());
        a.receive(&mut ours, opening).unwrap();
        assert!(a.generate(&mut ours, now, bytes).is_some());
        a.withdraw_for_delta();
        a.send_succeeded(a.id());

        // b loads the delta and tells a its heads, which are a's now.
        theirs.merge(&mut ours.clone()).unwrap();
        assert!(
            !b.have_all_theirs(&mut theirs),
            "b hasn't heard a's heads yet"
        );
        let after_delta = b.generate(&mut theirs, now, bytes).unwrap();
        b.send_succeeded(b.id());
        a.receive(&mut ours, after_delta).unwrap();

        let answer = a
            .generate(&mut ours, now, bytes)
            .expect("a answers with its heads");
        b.receive(&mut theirs, answer).unwrap();
        assert!(b.have_all_theirs(&mut theirs));
    }
}