//! Automerge synchronization: the sync worker, delivery with retries, inbound sync
//! messages, and the change publisher that feeds incremental streams.

use std::collections::HashSet;
use std::sync::Arc;
use std::sync::atomic::Ordering;

use anyhow::{Context, Result, anyhow, bail};
use automerge::sync::SyncDoc;
use tokio::sync::mpsc;
use tokio::time::Instant;
use tracing::{debug, info, trace, warn};

use crate::state::{AppState, Liveness, Sending, Standing};
use crate::swarm::peer::{peer_swarm, send_hello_once};
use crate::swarm::source;
use crate::swarm::throughput::SyncBytes;
use crate::transport::{self, ParsedMessage};

#[derive(Debug)]
pub(crate) enum SyncEvent {
    PeerValidated(String),
    DocumentChanged,
    MessageReceived(String),
    PeerNeedsSync(String),
}

/// Store paths touched by one batch of document changes.
#[derive(Debug)]
pub(crate) struct ChangeSet {
    pub(crate) paths: Vec<Vec<String>>,
}

async fn deliver_sync_to_peer(state: &AppState, peer: &str) -> Result<()> {
    if state.remake.fetching_snapshot() {
        // Rejoining: the snapshot comes as a download, not in a sync message. Sync
        // starts once it is here (`remake::catchup`).
        trace!("Waiting for the snapshot download; not syncing yet");
        return Ok(());
    }
    let Some(peer_state_mutex) = state.get_peer_state(peer).await else {
        trace!("Peer was forgotten; skipping synchronization");
        return Ok(());
    };

    let we_ready = state.is_ready().await;

    let validated_swarm = {
        let peer_state = peer_state_mutex.lock().await;

        if matches!(peer_state.liveness, Liveness::Dormant { .. }) {
            // Reactivation re-handshakes, and PeerValidated then restarts sync.
            trace!("Skipping synchronization with dormant peer");
            return Ok(());
        }

        if peer_state
            .epoch
            .is_some_and(|epoch| epoch != state.local_epoch())
        {
            // Unrelated document histories must never be synced.
            trace!("Skipping synchronization with peer on another epoch");
            return Ok(());
        }

        let Standing::Member { swarm_id } = &peer_state.standing else {
            return Ok(());
        };

        if let Some(wait) = source::defer(state, peer, &peer_state) {
            trace!(
                wait_ms = wait.as_millis(),
                "Catching up from another peer first; not syncing with this one yet"
            );
            // Done earlier when the initial sync completes (every peer is enqueued
            // then); this only matters if the source goes quiet.
            let tx = state.sync_events.clone();
            let peer = peer.to_owned();
            tokio::spawn(async move {
                tokio::time::sleep(wait).await;
                let _ = tx.send(SyncEvent::PeerNeedsSync(peer));
            });
            return Ok(());
        }

        if peer_state
            .sync
            .leave_opening(we_ready, peer_state.reports_ready)
        {
            // Its own opening (or anything else from it) brings us back here.
            trace!("Leaving the opening sync message to the peer, which isn't ready");
            return Ok(());
        }

        if peer_state.sync.state.their_heads.is_none()
            && !crate::swarm::topology::is_neighbour(state, peer)
        {
            // We only answer peers that aren't our neighbours.
            trace!("Not opening sync with a peer that isn't our neighbour");
            return Ok(());
        }

        swarm_id.clone()
    };

    let swarm_id = state
        .swarm_id()
        .await
        .ok_or_else(|| anyhow!("Cannot sync without an established swarm ID"))?;

    if validated_swarm != swarm_id {
        bail!("Peer {peer} belongs to a different swarm");
    }

    {
        let mut peer_state = peer_state_mutex.lock().await;
        if let Some(sending) = &mut peer_state.sync.sending {
            sending.regenerate = true;
            return Ok(());
        }
        if peer_state
            .sync
            .retry_at
            .is_some_and(|at| Instant::now() < at)
        {
            // The retry timer enqueues the peer once it's due.
            return Ok(());
        }
    }

    let encoded = {
        let doc = Arc::clone(&state.doc);
        let psm = Arc::clone(&peer_state_mutex);
        let transport = Arc::clone(&state.transport);
        let peer_owned = peer.to_owned();
        let ready = state.is_ready().await;
        // Before this message is counted.
        let bytes_before = SyncBytes::of(&state.traffic, peer);

        let generated =
            tokio::task::spawn_blocking(move || -> Result<Option<transport::EncodedMessage>> {
                // Label the message with the epoch of the document it was generated
                // from, read under the same lock.
                let (epoch, message) = {
                    let mut replica = doc.blocking_write();
                    let mut peer_state = psm.blocking_lock();
                    let epoch = replica.epoch;
                    let sent_before = peer_state.sync.state.sent_hashes.clone();
                    let message = {
                        let sync_doc = replica.doc.sync();
                        sync_doc.generate_sync_message(&mut peer_state.sync.state)
                    };
                    // A peer whose heads are ours already has everything; Automerge
                    // answers only when it has something to say, so it will stay
                    // silent. Waiting for an answer would only end in a pointless
                    // resend.
                    let in_step = peer_state.sync.state.their_heads.as_ref()
                        == Some(&replica.doc.get_heads());
                    let sent: Vec<_> = peer_state
                        .sync
                        .state
                        .sent_hashes
                        .difference(&sent_before)
                        .copied()
                        .collect();
                    if in_step {
                        peer_state.sync.reply_settled();
                    } else if message.is_some() {
                        peer_state
                            .sync
                            .expect_reply(sent.clone(), Instant::now(), bytes_before);
                    }
                    if message.is_some() {
                        // Marked under the same lock as the generation, so no second
                        // message is generated while this one is handed over.
                        peer_state.sync.sending = Some(Sending {
                            sent,
                            regenerate: false,
                        });
                    }
                    (epoch, message)
                };
                let Some(message) = message else {
                    return Ok(None);
                };
                let envelope = transport::SyncEnvelope {
                    swarm_id,
                    epoch,
                    message,
                    ready,
                };
                transport
                    .encode(&peer_owned, &ParsedMessage::Sync(envelope))
                    .map(Some)
            })
            .await
            .context("Sync generation task panicked")?;
        let generated = match generated {
            Ok(generated) => generated,
            Err(error) => {
                // Not encodable, so nothing went out: take back what it would have carried.
                let _ = peer_state_mutex
                    .lock()
                    .await
                    .sync
                    .send_failed(Instant::now());
                return Err(error);
            }
        };

        let Some(encoded) = generated else {
            return Ok(());
        };

        encoded
    };

    // No state locks are held across the send.
    match state.transport.send_encoded(&encoded).await {
        Ok(()) => {
            let regenerate = peer_state_mutex.lock().await.sync.send_succeeded();

            trace!(regenerate, "Outbound sync message accepted by Mycelium");

            if regenerate {
                state
                    .sync_events
                    .send(SyncEvent::PeerNeedsSync(peer.to_owned()))
                    .map_err(|error| anyhow!("Failed to notify sync worker: {error}"))?;
            }

            Ok(())
        }
        Err(error) => {
            let Some(delay) = peer_state_mutex
                .lock()
                .await
                .sync
                .send_failed(Instant::now())
            else {
                debug!("Sync message refused, but a reset already discarded it");
                return Ok(());
            };

            warn!(
                retry_ms = delay.as_millis(),
                error = %format!("{error:#}"),
                "Sync delivery failed; generating a fresh message later"
            );

            let tx = state.sync_events.clone();
            let peer = peer.to_owned();
            tokio::spawn(async move {
                tokio::time::sleep(delay).await;
                let _ = tx.send(SyncEvent::PeerNeedsSync(peer));
            });

            Ok(())
        }
    }
}

pub(crate) async fn handle_sync(
    state: &AppState,
    peer: &str,
    envelope: transport::SyncEnvelope,
) -> Result<()> {
    debug!(peer = %peer, "Processing sync envelope");
    // This message is counted already (`transport::traffic`).
    let arrived = (Instant::now(), SyncBytes::of(&state.traffic, peer));

    let peer_state_mutex = state.get_or_create_peer_state(peer).await;

    let needs_hello = {
        let peer_state = peer_state_mutex.lock().await;

        match &peer_state.standing {
            Standing::Stranger | Standing::Admitted => {
                // Not processed: our handshake reply resets the session anyway, and the
                // sender's reply timeout resends what still matters.
                debug!(peer = %peer, "Dropping sync message from a peer that isn't a member yet");
                true
            }
            Standing::Member { swarm_id } if swarm_id != &envelope.swarm_id => {
                bail!(
                    "Sync swarm mismatch for peer {peer}: validated_swarm={swarm_id}, envelope_swarm={}",
                    envelope.swarm_id
                );
            }
            Standing::Member { .. } => false,
        }
    };

    if needs_hello {
        // The peer treats us as a member but we don't know it, typically because we
        // restarted. Our HELLO makes it reset its sync state and reply.
        if let Err(error) = send_hello_once(state, peer).await {
            warn!(peer = %peer, error = %format!("{error:#}"), "HELLO to unknown sync sender failed; will retry");
        }
        return Ok(());
    }

    let expected_swarm = state
        .swarm_id()
        .await
        .ok_or_else(|| anyhow!("Missing expected swarm ID after peer membership"))?;
    if envelope.swarm_id != expected_swarm {
        bail!(
            "Incompatible swarm from {peer}: expected {}, received {}",
            expected_swarm,
            envelope.swarm_id
        );
    }

    let epoch_label = envelope.epoch;
    let from_ready = envelope.ready;

    let applied = {
        let doc = Arc::clone(&state.doc);
        let peer_state_mutex = Arc::clone(&peer_state_mutex);
        let stats = Arc::clone(&state.stats);
        let throughput = Arc::clone(&state.throughput);
        let rtt = state
            .lock_rtt()
            .effective(peer, Instant::now(), state.timing.rtt_fresh_for());
        let check_caught_up = !state.initial_sync_done.load(Ordering::Acquire);
        // While catching up, everything arrives late; that says nothing about propagation.
        let measure_propagation = !check_caught_up;
        let epoch = envelope.epoch;
        let message = envelope.message;
        let peer = peer.to_owned();

        // Some((caught_up, heads changed)) when applied, None when the message belongs
        // to another epoch.
        tokio::task::spawn_blocking(move || -> Result<Option<(bool, bool)>> {
            let mut replica = doc.blocking_write();

            // Authoritative check, under the same lock as the apply: the document may
            // have switched epochs since the message was accepted.
            if replica.epoch != epoch {
                return Ok(None);
            }

            let mut peer_state = peer_state_mutex.blocking_lock();
            peer_state.reports_ready = Some(from_ready);
            let heads_before = replica.doc.get_heads();

            replica
                .doc
                .sync()
                .receive_sync_message(&mut peer_state.sync.state, message)
                .context("Automerge synchronization failed")?;
            peer_state.last_sync_rx = Some(Instant::now());
            if let Some(wait) = &peer_state.sync.reply {
                throughput
                    .lock()
                    .unwrap_or_else(|poisoned| poisoned.into_inner())
                    .exchange(&peer, wait.started, arrived, rtt);
            }
            peer_state.sync.reply_settled();

            resend_if_needed(&peer, &mut peer_state.sync.state);
            let heads_after = replica.doc.get_heads();
            stats.heads_seen(heads_after.len());
            let changed = heads_after != heads_before;
            if changed && measure_propagation {
                let now_ms = crate::state::unix_millis() as i64;
                // Time 0: not stamped (snapshots, older nodes).
                let delays: Vec<i64> = replica
                    .doc
                    .get_changes(&heads_before)
                    .iter()
                    .map(|change| change.timestamp())
                    .filter(|&time| time > 0)
                    .map(|time| now_ms - time)
                    .collect();
                stats.propagated(delays);
            }

            Ok(Some((
                // Only from a ready peer: a peer still catching up has empty heads, which
                // every document "contains".
                check_caught_up
                    && from_ready
                    && peer_state
                        .sync
                        .state
                        .their_heads
                        .as_ref()
                        .is_some_and(|heads| {
                            heads
                                .iter()
                                .all(|h| replica.doc.get_change_by_hash(h).is_some())
                        }),
                changed,
            )))
        })
        .await
        .context("Sync task panicked")??
    };

    let Some((caught_up, changed)) = applied else {
        debug!(peer = %peer, epoch = %epoch_label, "Dropped sync message from another epoch");
        return Ok(());
    };

    source::reported(state, peer, from_ready, None);

    if caught_up && !state.initial_sync_done.swap(true, Ordering::AcqRel) {
        info!(peer = %peer, "Initial sync complete");
        state
            .stats
            .caught_up(state.traffic.total().received.sync.bytes);
        source::finished(state, peer);
    }

    state.notify_store_changed();

    crate::swarm::relay::pulled(state, peer);
    if changed {
        crate::swarm::relay::announce(state, peer).await;
    }

    state
        .sync_events
        .send(SyncEvent::MessageReceived(peer.to_owned()))
        .map_err(|error| anyhow!("Failed to notify sync worker: {error}"))?;

    Ok(())
}

/// 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.)
fn resend_if_needed(peer: &str, sync_state: &mut automerge::sync::State) {
    let Some(need) = &sync_state.their_need else {
        return;
    };
    if need
        .iter()
        .any(|hash| sync_state.sent_hashes.contains(hash))
    {
        debug!(
            peer = %peer,
            sent = sync_state.sent_hashes.len(),
            "Peer needs changes we had sent; forgetting what was sent"
        );
        sync_state.sent_hashes.clear();
    }
}

pub(crate) async fn run_sync_worker(state: AppState, mut rx: mpsc::UnboundedReceiver<SyncEvent>) {
    let mut dirty_peers = HashSet::<String>::new();
    let mut document_dirty = false;

    info!("Starting synchronization worker");

    loop {
        if dirty_peers.is_empty() && !document_dirty {
            match rx.recv().await {
                Some(event) => match event {
                    SyncEvent::PeerValidated(peer)
                    | SyncEvent::MessageReceived(peer)
                    | SyncEvent::PeerNeedsSync(peer) => {
                        dirty_peers.insert(peer);
                    }
                    SyncEvent::DocumentChanged => {
                        document_dirty = true;
                    }
                },
                None => {
                    warn!("Synchronization event channel closed");
                    return;
                }
            }
        }

        // Coalesce event bursts.
        while let Ok(event) = rx.try_recv() {
            match event {
                SyncEvent::PeerValidated(peer)
                | SyncEvent::MessageReceived(peer)
                | SyncEvent::PeerNeedsSync(peer) => {
                    dirty_peers.insert(peer);
                }
                SyncEvent::DocumentChanged => {
                    document_dirty = true;
                }
            }
        }

        // Let pending HTTP task locks wake up and complete their lifecycle transitions gracefully.
        tokio::task::yield_now().await;

        if document_dirty {
            document_dirty = false;

            // Lock and extract only the cloned string keys into a small transient vector
            let peer_keys: Vec<String> = {
                let peers_guard = state.peer_states.read().await;
                peers_guard.keys().cloned().collect()
            };

            // Our own changes go to our neighbours; they pass them on. Others only get
            // answers to what they send us.
            for peer in peer_keys {
                if peer_swarm(&state, &peer).await.is_some()
                    && crate::swarm::topology::is_neighbour(&state, &peer)
                {
                    dirty_peers.insert(peer);
                }
            }
        }

        let peers: Vec<String> = dirty_peers.drain().collect();

        for peer in peers {
            if let Err(error) = deliver_sync_to_peer(&state, &peer).await {
                warn!(
                    peer = %peer,
                    error = %format!("{error:#}"),
                    "Synchronization processing failed"
                );
            }
        }
    }
}

// --- change publisher --------------------------------------------------------

fn patch_store_path(patch: &automerge::Patch) -> Vec<String> {
    let mut path: Vec<String> = patch
        .path
        .iter()
        .map(|(_, prop)| match prop {
            automerge::Prop::Map(key) => key.clone(),
            automerge::Prop::Seq(index) => index.to_string(),
        })
        .collect();

    match &patch.action {
        automerge::PatchAction::PutMap { key, .. }
        | automerge::PatchAction::DeleteMap { key, .. } => path.push(key.clone()),
        // List, text and counter changes: the containing object changed.
        _ => {}
    }
    path
}

pub(crate) fn paths_related(a: &[String], b: &[String]) -> bool {
    a.starts_with(b) || b.starts_with(a)
}

pub(crate) async fn run_change_publisher(state: AppState) {
    let mut changed = state.change_tx.subscribe();
    let mut last_heads = state.write_doc(|doc| doc.get_heads()).await;

    info!("Starting change publisher");

    while changed.changed().await.is_ok() {
        let before = last_heads.clone();

        let (heads, paths, conflicts) = state
            .write_doc(move |doc| {
                let heads = doc.get_heads();
                if heads == before {
                    return (heads, Vec::new(), Vec::new());
                }

                // After an epoch switch the previous heads belong to the old document:
                // report the root as changed, which every stream is related to.
                if !before.iter().all(|h| doc.get_change_by_hash(h).is_some()) {
                    return (heads, vec![Vec::new()], Vec::new());
                }

                let mut paths = Vec::new();
                let mut conflicts = Vec::new();
                for patch in doc.diff(&before, &heads) {
                    let path = patch_store_path(&patch);
                    if matches!(
                        patch.action,
                        automerge::PatchAction::PutMap { conflict: true, .. }
                    ) {
                        conflicts.push(path.clone());
                    }
                    paths.push(path);
                }
                (heads, paths, conflicts)
            })
            .await;

        if heads == last_heads {
            continue;
        }
        last_heads = heads;

        for path in &conflicts {
            warn!(path = ?path, "Concurrent writes produced a conflict");
        }

        // Fails only when nobody is subscribed.
        let _ = state.change_events.send(Arc::new(ChangeSet { paths }));
    }
}