//! 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::atomic::Ordering;
use std::sync::Arc;
use std::time::Duration;

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

use crate::peer::{peer_cluster, send_hello_once};
use crate::state::{AppState, Liveness, PeerMembership, PeerState, PendingOutbound};
use crate::transport::{self, ParsedMessage};

const INITIAL_SYNC_RETRY: Duration = Duration::from_millis(500);
const MAX_SYNC_RETRY: Duration = Duration::from_secs(30);

#[derive(Debug)]
pub(crate) enum SyncEvent {
    PeerValidated(String),
    DocumentChanged,
    MessageReceived(String),
    PeerNeedsSync(String),
    RetryDelivery(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, retry_delivery: bool) -> Result<()> {
    let Some(peer_state_mutex) = state.get_peer_state(peer).await else {
        trace!("Peer was forgotten; skipping synchronization");
        return Ok(());
    };

    let validated_cluster = {
        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(());
        }

        match &peer_state.membership {
            PeerMembership::Unknown => return Ok(()),
            PeerMembership::Member { cluster_id } => cluster_id.clone(),
        }
    };

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

    if validated_cluster != cluster_id {
        bail!("Peer {peer} belongs to a different cluster");
    }

    let existing = {
        let mut guard = peer_state_mutex.lock().await;
        let peer_state = &mut *guard;
        match &peer_state.pending_outbound_sync {
            Some(entry) if retry_delivery => Some((entry.message.clone(), entry.next_attempt)),
            Some(_) => {
                peer_state.needs_generate = true;
                return Ok(());
            }
            None => None,
        }
    };

    let encoded = match existing {
        Some((message, next_attempt)) => {
            if Instant::now() < next_attempt {
                return Ok(());
            }
            message
        }
        None if retry_delivery => return Ok(()),
        None => {
            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 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();
                        peer_state.needs_generate = false;
                        let epoch = replica.epoch;
                        let sync_doc = replica.doc.sync();
                        (
                            epoch,
                            sync_doc.generate_sync_message(&mut peer_state.sync_state),
                        )
                    };
                    let Some(message) = message else {
                        return Ok(None);
                    };
                    let envelope = transport::SyncEnvelope {
                        cluster_id,
                        epoch,
                        message,
                    };
                    transport
                        .encode(&peer_owned, &ParsedMessage::Sync(envelope))
                        .map(Some)
                },
            )
            .await
            .context("Sync generation task panicked")??;

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

            peer_state_mutex.lock().await.pending_outbound_sync = Some(PendingOutbound {
                message: encoded.clone(),
                retry_delay: INITIAL_SYNC_RETRY,
                next_attempt: Instant::now(),
            });

            encoded
        }
    };

    // No state locks are held across the send.
    match state.transport.send_encoded(&encoded).await {
        Ok(()) => {
            let regenerate = {
                let mut peer_state = peer_state_mutex.lock().await;
                peer_state.pending_outbound_sync = None;
                std::mem::take(&mut peer_state.needs_generate)
            };

            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 delay = {
                let mut peer_state = peer_state_mutex.lock().await;

                let Some(ref mut entry) = peer_state.pending_outbound_sync else {
                    debug!("Pending envelope discarded by a sync reset during delivery");
                    return Ok(());
                };

                entry.retry_delay = entry.retry_delay.saturating_mul(2).min(MAX_SYNC_RETRY);
                entry.next_attempt = Instant::now() + entry.retry_delay;
                entry.retry_delay
            };

            warn!(
                retry_ms = delay.as_millis(),
                error = %format!("{error:#}"),
                "Sync delivery failed; retaining the same outbound envelope"
            );

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

            Ok(())
        }
    }
}

pub(crate) async fn handle_sync(
    state: &AppState,
    peer: &str,
    envelope: transport::SyncEnvelope,
) -> Result<()> {
    debug!(peer = %peer, "Processing sync envelope");

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

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

        match &peer_state.membership {
            PeerMembership::Unknown => {
                let duplicate = peer_state.pending_inbound_sync.iter().any(|pending| {
                    pending.message == envelope.message
                        && pending.cluster_id == envelope.cluster_id
                        && pending.epoch == envelope.epoch
                });
                if !duplicate {
                    peer_state.pending_inbound_sync.push(envelope.clone());
                    info!(
                        peer = %peer,
                        pending_count = peer_state.pending_inbound_sync.len(),
                        "Queued sync envelope pending peer membership"
                    );
                } else {
                    trace!(peer = %peer, "Ignoring duplicate pending sync envelope");
                }
                true
            }
            PeerMembership::Member { cluster_id } if cluster_id != &envelope.cluster_id => {
                bail!(
                    "Sync cluster mismatch for peer {peer}: validated_cluster={cluster_id}, envelope_cluster={}",
                    envelope.cluster_id
                );
            }
            PeerMembership::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_cluster = state
        .cluster_id()
        .await
        .ok_or_else(|| anyhow!("Missing expected cluster ID after peer membership"))?;
    if envelope.cluster_id != expected_cluster {
        bail!(
            "Incompatible cluster from {peer}: expected {}, received {}",
            expected_cluster,
            envelope.cluster_id
        );
    }

    let epoch_label = envelope.epoch;

    let applied = {
        let doc = Arc::clone(&state.doc);
        let peer_state_mutex = Arc::clone(&peer_state_mutex);
        let check_caught_up = !state.initial_sync_done.load(Ordering::Acquire);
        let epoch = envelope.epoch;
        let message = envelope.message;
        let peer = peer.to_owned();

        // Some(caught_up) when applied, None when the message belongs to another epoch.
        tokio::task::spawn_blocking(move || -> Result<Option<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();

            let heads_before = replica.doc.get_heads();
            let their_heads_before = peer_state.sync_state.their_heads.clone();

            replica
                .doc
                .sync()
                .receive_sync_message(&mut peer_state.sync_state, message)
                .context("Automerge synchronization failed")?;
            peer_state.last_sync_rx = Some(Instant::now());

            let heads_after = replica.doc.get_heads();
            let unproductive = heads_before == heads_after
                && peer_state.sync_state.their_heads == their_heads_before
                && peer_state.sync_state.their_heads.as_ref() != Some(&heads_after);
            break_sync_stall(&peer, &mut peer_state, unproductive);

            Ok(Some(
                check_caught_up
                    && 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())
                        }),
            ))
        })
        .await
        .context("Sync task panicked")??
    };

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

    if caught_up && !state.initial_sync_done.swap(true, Ordering::AcqRel) {
        info!(peer = %peer, "Initial sync complete");
    }

    state.notify_store_changed();

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

    Ok(())
}

/// Exchanges in a row that changed nothing on either side before `sent_hashes` is
/// cleared, and before a full two-sided reset (re-handshake).
const STALL_RESEND_NEEDED: u32 = 2;
const STALL_CLEAR_SENT: u32 = 8;
const STALL_RESET: u32 = 32;

/// Breaks sync livelocks.
///
/// Automerge never resends a change it believes the peer has received (`sent_hashes`).
/// If that change was lost on the way (a dropped or undeliverable message, an epoch
/// check, a reset of only one side's state), the two sides keep exchanging messages
/// with differing heads and no changes, as fast as the transport allows. The keepalive
/// heal never fires then, because sync traffic keeps flowing.
///
/// * A change the peer keeps asking for is sent again.
/// * After a few unproductive exchanges, forget what we sent: the next message carries
///   everything the peer's bloom filter lacks.
/// * If that doesn't help either, reset both sides through a fresh handshake.
fn break_sync_stall(peer: &str, peer_state: &mut PeerState, unproductive: bool) {
    if !unproductive {
        peer_state.stalled_exchanges = 0;
        return;
    }
    peer_state.stalled_exchanges += 1;

    let sync_state = &mut peer_state.sync_state;

    // Not on the first unproductive exchange: when messages cross, the peer asks for a
    // change that is merely still in flight, and resending it would only duplicate it.
    if peer_state.stalled_exchanges >= STALL_RESEND_NEEDED {
        if let Some(need) = &sync_state.their_need {
            let before = sync_state.sent_hashes.len();
            sync_state.sent_hashes.retain(|hash| !need.contains(hash));
            let resend = before - sync_state.sent_hashes.len();
            if resend > 0 {
                debug!(peer = %peer, resend, "Peer still needs changes we had sent; sending them again");
            }
        }
    }

    match peer_state.stalled_exchanges {
        STALL_CLEAR_SENT => {
            warn!(
                peer = %peer,
                sent = sync_state.sent_hashes.len(),
                their_heads = ?sync_state.their_heads,
                shared_heads = ?sync_state.shared_heads,
                "Sync is stalled; forgetting which changes were sent"
            );
            sync_state.sent_hashes.clear();
        }
        STALL_RESET => {
            warn!(peer = %peer, "Sync is still stalled; resetting it on both sides");
            peer_state.stalled_exchanges = 0;
            peer_state.request_resync();
        }
        _ => {}
    }
}

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

    info!("Starting synchronization worker");

    loop {
        if dirty_peers.is_empty() && retry_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::RetryDelivery(peer) => {
                        retry_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::RetryDelivery(peer) => {
                    retry_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()
            };

            for peer in peer_keys {
                if peer_cluster(&state, &peer).await.is_some() {
                    dirty_peers.insert(peer);
                }
            }
        }

        // Retry existing envelopes separately from generating new ones.
        let retries: Vec<String> = retry_peers.drain().collect();

        for peer in retries {
            if let Err(error) = deliver_sync_to_peer(&state, &peer, true).await {
                warn!(
                    peer = %peer,
                    error = %format!("{error:#}"),
                    "Sync delivery retry failed"
                );
            }
        }

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

        for peer in peers {
            if let Err(error) = deliver_sync_to_peer(&state, &peer, false).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 }));
    }
}