//! Peers: handshake (HELLO / HELLO_REPLY), liveness and keepalives, self-healing,
//! cluster-wide forget records, and dispatch of inbound messages.

use std::sync::Arc;
use std::time::Duration;

use anyhow::{anyhow, bail, Context, Result};
use automerge::AutoCommit;
use sha2::{Digest, Sha256};
use tokio::sync::Mutex;
use tokio::time::{Instant, MissedTickBehavior};
use tracing::{debug, info, instrument, trace, warn};

use crate::replica::{Epoch, EpochRelation};
use crate::state::{AppState, HelloState, Liveness, PeerMembership, PeerState};
use crate::store::{self, StoreError};
use crate::sync::{handle_sync, SyncEvent};
use crate::transport::{self, ParsedMessage};

const HELLO_RETRY_INITIAL: Duration = Duration::from_secs(1);
const HELLO_RETRY_MAX: Duration = Duration::from_secs(30);

const PEER_TICK: Duration = Duration::from_millis(500);

const META_ROOT: &str = "_continuum";
const FORGET_CHECK_INTERVAL: Duration = Duration::from_secs(5);

// --- forget records ----------------------------------------------------------

fn forgotten_dir() -> Vec<String> {
    vec![META_ROOT.to_owned(), "forgotten".to_owned()]
}

pub(crate) fn forgotten_path(pk: &str) -> Vec<String> {
    let mut path = forgotten_dir();
    path.push(pk.to_owned());
    path
}

fn forgotten_peers(doc: &AutoCommit) -> Vec<String> {
    store::keys_at(doc, &forgotten_dir()).unwrap_or_default()
}

async fn apply_forget_records(state: &AppState) {
    // Before the initial sync we can't judge the records, and we have nothing to delete.
    if !state.is_ready().await {
        return;
    }

    let forgotten = forgotten_peers(&state.doc.read().await.doc);
    if forgotten.is_empty() {
        return;
    }

    let now = Instant::now();
    let mut cancel = Vec::new();

    for pk in forgotten {
        if pk == state.local_id() {
            cancel.push(pk); // we are evidently alive
            continue;
        }

        let Some(peer_state_mutex) = state.get_peer_state(&pk).await else {
            continue; // already forgotten here
        };

        let alive = peer_state_mutex.lock().await.is_alive(&state.timing, now);

        if alive {
            cancel.push(pk);
        } else if state.peer_states.write().await.remove(&pk).is_some() {
            info!(peer = %pk, "Forgetting peer per cluster-wide record");
        }
    }

    if cancel.is_empty() {
        return;
    }

    // Individual failures are logged inside; NotFound means another node already
    // cancelled the record.
    let _ = state
        .mutate_doc(move |doc| {
            for pk in &cancel {
                match store::delete_path(&mut *doc, &forgotten_path(pk), false) {
                    Ok(()) => info!(peer = %pk, "Peer is alive; cancelling cluster-wide forget"),
                    Err(StoreError::NotFound) => {}
                    Err(error) => {
                        warn!(peer = %pk, error = %error, "Failed to cancel forget record")
                    }
                }
            }
            Ok::<(), StoreError>(())
        })
        .await;
}

// --- keepalive and self-healing ----------------------------------------------

fn heads_digest(doc: &mut AutoCommit) -> transport::HeadsDigest {
    let mut heads = doc.get_heads();
    heads.sort_unstable();
    let mut hasher = Sha256::new();
    for head in &heads {
        hasher.update(head.0);
    }
    hasher.finalize().into()
}

pub(crate) async fn local_keepalive(state: &AppState) -> transport::Keepalive {
    // AutoCommit::get_heads takes &mut self (it closes the pending transaction).
    // Digest and epoch come from the same locked replica, so they always match.
    let mut replica = state.doc.write().await;
    transport::Keepalive {
        heads_digest: heads_digest(&mut replica.doc),
        epoch: replica.epoch,
    }
}

async fn heal_if_diverged(
    state: &AppState,
    peer: &str,
    local: &transport::Keepalive,
    remote: &transport::Keepalive,
) {
    // Digests of documents from different epochs are never equal and never converge.
    if local.epoch != remote.epoch {
        return;
    }

    let Some(peer_state_mutex) = state.get_peer_state(peer).await else {
        return;
    };

    let heal = peer_state_mutex.lock().await.needs_heal(
        &local.heads_digest,
        &remote.heads_digest,
        &state.timing,
        Instant::now(),
    );

    if heal {
        warn!(peer = %peer, "Documents differ and sync is idle; re-handshaking to restart sync");
        ensure_handshake(state, peer).await;
    }
}

// --- peer discovery and HELLO ------------------------------------------------

pub(crate) async fn remember_peer(state: &AppState, peer: &str) {
    if peer.is_empty() {
        warn!("Ignoring empty peer public key");
        return;
    }

    if peer == state.local_id() {
        debug!("Ignoring own public key");
        return;
    }

    let mut peer_states = state.peer_states.write().await;

    if !peer_states.contains_key(peer) {
        peer_states.entry(peer.to_owned()).or_default();
        info!(
            peer = %peer,
            peer_count = peer_states.len(),
            "Discovered peer"
        );
    } else {
        trace!(peer = %peer, "Peer already registered");
    }
}

pub(crate) async fn send_hello_once(state: &AppState, peer: &str) -> Result<()> {
    let Some(peer_state_mutex) = state.get_peer_state(peer).await else {
        trace!(peer = %peer, "Peer was forgotten; not sending HELLO");
        return Ok(());
    };

    let claimed = {
        let mut peer_state = peer_state_mutex.lock().await;
        if peer_state.hello_state == HelloState::NotSent {
            peer_state.hello_state = HelloState::Sending {
                next_delay: HELLO_RETRY_INITIAL,
            };
            true
        } else {
            trace!(peer = %peer, hello_state = ?peer_state.hello_state, "not sending");
            false
        }
    };

    if !claimed {
        return Ok(());
    }

    debug!(to = %peer, "Sending HELLO");
    transmit_hello(state, peer, HELLO_RETRY_INITIAL).await
}

async fn retry_hello_if_due(state: &AppState, peer: &str) -> Result<()> {
    let Some(peer_state_mutex) = state.get_peer_state(peer).await else {
        return Ok(());
    };

    let delay = {
        let mut peer_state = peer_state_mutex.lock().await;
        match peer_state.hello_state {
            HelloState::AwaitingReply {
                next_attempt,
                next_delay,
            } if Instant::now() >= next_attempt => {
                peer_state.hello_state = HelloState::Sending { next_delay };
                next_delay
            }
            _ => return Ok(()),
        }
    };

    debug!(
        peer = %peer,
        next_retry_ms = delay.as_millis(),
        "No HELLO_REPLY yet; re-sending HELLO"
    );
    transmit_hello(state, peer, delay).await
}

/// Sends a HELLO without holding any lock, then schedules the next retry.
/// Success and failure both end in AwaitingReply: Mycelium accepting the message
/// says nothing about whether the peer received it.
async fn transmit_hello(state: &AppState, peer: &str, delay: Duration) -> Result<()> {
    let result = state
        .transport
        .send_message(
            peer,
            &ParsedMessage::Hello(transport::HelloRequest {
                epoch: state.local_epoch(),
            }),
        )
        .await;

    if let Some(peer_state_mutex) = state.get_peer_state(peer).await {
        let mut peer_state = peer_state_mutex.lock().await;

        if matches!(peer_state.hello_state, HelloState::Sending { .. }) {
            let (wait, next_delay) = match peer_state.liveness {
                Liveness::Dormant { .. } => {
                    (state.timing.dormant_probe, state.timing.dormant_probe)
                }
                Liveness::Active => (delay, delay.saturating_mul(2).min(HELLO_RETRY_MAX)),
            };
            peer_state.hello_state = HelloState::AwaitingReply {
                next_attempt: Instant::now() + wait,
                next_delay,
            };
        }
    }

    result.context("Sending HELLO")
}

async fn ensure_handshake(state: &AppState, peer: &str) {
    if let Err(error) = send_hello_once(state, peer).await {
        warn!(peer = %peer, error = %format!("{error:#}"), "HELLO failed; will retry");
    }
}

/// Sends a HELLO_REPLY if one is owed and our cluster ID is known.
/// Otherwise it stays owed and the peer worker tries again.
async fn send_owed_hello_reply(state: &AppState, peer: &str) {
    if state.cluster_id().await.is_none() {
        return;
    }

    let Some(peer_state_mutex) = state.get_peer_state(peer).await else {
        return;
    };

    let owed = std::mem::take(&mut peer_state_mutex.lock().await.reply_owed);
    if !owed {
        return;
    }

    if let Err(error) = send_hello_reply(state, peer).await {
        warn!(peer = %peer, error = %format!("{error:#}"), "HELLO_REPLY failed; will retry");
        peer_state_mutex.lock().await.reply_owed = true;
    }
}

async fn flush_owed_hello_replies(state: &AppState) {
    let peers: Vec<String> = state.peer_states.read().await.keys().cloned().collect();
    for peer in &peers {
        send_owed_hello_reply(state, peer).await;
    }
}

async fn send_hello_reply(state: &AppState, peer: &str) -> Result<()> {
    let cluster_id = state
        .cluster_id()
        .await
        .ok_or_else(|| anyhow!("Cannot send HELLO_REPLY without an established cluster ID"))?;

    let candidates: Vec<(String, Arc<Mutex<PeerState>>)> = state
        .peer_states
        .read()
        .await
        .iter()
        .filter(|(known, _)| known.as_str() != peer)
        .map(|(known, m)| (known.clone(), Arc::clone(m)))
        .collect();

    let mut peers = Vec::new();
    for (known, peer_state_mutex) in candidates {
        let peer_state = peer_state_mutex.lock().await;
        if peer_state.liveness == Liveness::Active
            && peer_state.hello_state == HelloState::Established
        {
            peers.push(known);
        }
    }

    let reply = transport::HelloReply {
        cluster_id,
        peers,
        epoch: state.local_epoch(),
    };

    let ready = state.is_ready().await;

    debug!(to = %peer, ready, "Sending HELLO_REPLY");

    state
        .transport
        .send_message(peer, &ParsedMessage::HelloReply(reply))
        .await
        .context("Sending HELLO_REPLY")?;

    trace!(to = %peer, "HELLO_REPLY sent");
    Ok(())
}

async fn discover_peers(state: &AppState, from_peer: &str, advertised_peers: &[String]) {
    // TODO here we could add a trace t oshwo advertised_peers
    for discovered_peer in advertised_peers {
        if discovered_peer == from_peer {
            continue;
        }

        remember_peer(state, discovered_peer).await;

        if let Err(error) = send_hello_once(state, discovered_peer).await {
            warn!(
                peer = %discovered_peer,
                error = %format!("{error:#}"),
                "Failed to initiate HELLO with discovered peer"
            );
        }
    }
}

// --- peer worker -------------------------------------------------------------

pub(crate) async fn run_peer_worker(state: AppState) {
    let mut tick = tokio::time::interval(PEER_TICK);
    tick.set_missed_tick_behavior(MissedTickBehavior::Delay);

    info!("Starting peer maintenance worker");

    let mut last_forget_check: Option<Instant> = None;

    loop {
        tick.tick().await;

        // While joining, keep handshakes fast; dormancy would slow probes to minutes.
        let ready = state.is_ready().await;

        let peers: Vec<(String, Arc<Mutex<PeerState>>)> = state
            .peer_states
            .read()
            .await
            .iter()
            .map(|(pk, m)| (pk.clone(), Arc::clone(m)))
            .collect();

        for (peer, peer_state_mutex) in &peers {
            let now = Instant::now();

            let (became_dormant, ping_due) = {
                let mut peer_state = peer_state_mutex.lock().await;
                let became_dormant = ready && peer_state.check_dormancy(&state.timing, now);
                let ping_due = peer_state.take_keepalive_due(&state.timing, now);
                (became_dormant, ping_due)
            };

            if became_dormant {
                warn!(
                    peer = %peer,
                    silent_secs = state.timing.dormant_after.as_secs(),
                    probe_secs = state.timing.dormant_probe.as_secs(),
                    "Peer went dormant"
                );
            }

            if ping_due {
                let local = local_keepalive(&state).await;
                if let Err(error) = state
                    .transport
                    .send_message(peer, &ParsedMessage::Ping(local))
                    .await
                {
                    debug!(peer = %peer, error = %format!("{error:#}"), "Keepalive PING failed");
                }
            }

            if let Err(error) = retry_hello_if_due(&state, peer).await {
                warn!(peer = %peer, error = %format!("{error:#}"), "HELLO retry failed");
            }

            send_owed_hello_reply(&state, peer).await;
        }

        if last_forget_check.is_none_or(|t| t.elapsed() >= FORGET_CHECK_INTERVAL) {
            apply_forget_records(&state).await;
            last_forget_check = Some(Instant::now());
        }
    }
}

// --- peer membership ---------------------------------------------------------

async fn establish_peer(
    state: &AppState,
    peer: &str,
    cluster_id: String,
) -> Result<Vec<transport::SyncEnvelope>> {
    let expected_cluster = state.cluster_id().await;

    if let Some(expected) = &expected_cluster {
        if expected != &cluster_id {
            bail!(
                "Incompatible cluster at {peer}: expected {}, received {}",
                expected,
                cluster_id
            );
        }
    }

    let peer_state_mutex = state.get_or_create_peer_state(peer).await;
    let mut peer_state = peer_state_mutex.lock().await;

    if let PeerMembership::Member {
        cluster_id: existing,
    } = &peer_state.membership
    {
        if existing != &cluster_id {
            bail!(
                "Peer {peer} changed cluster identity from {} to {}",
                existing,
                cluster_id
            );
        }
    }

    // Lock order: peer mutex -> lifecycle.
    if expected_cluster.is_none() {
        state.set_expected_cluster_id(cluster_id.clone()).await?;
    }

    if matches!(peer_state.membership, PeerMembership::Unknown) {
        peer_state.membership = PeerMembership::Member {
            cluster_id: cluster_id.clone(),
        };
    }

    // First reply completes the handshake: stop retries and sync from a clean slate.
    // Duplicate replies (caused by our own retries) don't reset again.
    if peer_state.hello_state != HelloState::Established {
        peer_state.hello_state = HelloState::Established;
        peer_state.reset_sync();
        info!(peer = %peer, "Handshake complete; sync state reset");
    }

    let pending = std::mem::take(&mut peer_state.pending_inbound_sync);

    info!(peer = %peer, cluster_id = %cluster_id, pending_messages = pending.len(), "Peer membership established");
    Ok(pending)
}

#[instrument(skip(state), fields(peer = %peer))]
pub(crate) async fn peer_cluster(state: &AppState, peer: &str) -> Option<String> {
    let peer_state_mutex = state.get_peer_state(peer).await?;
    let peer_state = peer_state_mutex.lock().await;
    match &peer_state.membership {
        PeerMembership::Unknown => None,
        PeerMembership::Member { cluster_id } => Some(cluster_id.clone()),
    }
}

// --- epochs ------------------------------------------------------------------

/// Records the epoch a peer advertised and compares it with ours. Returns true when both
/// are on the same epoch, the only case in which the documents may be synced.
async fn observe_peer_epoch(state: &AppState, peer: &str, remote: Epoch) -> bool {
    let local = state.local_epoch();
    let relation = local.relation_to(&remote);

    let Some(peer_state_mutex) = state.get_peer_state(peer).await else {
        return relation == EpochRelation::Same;
    };

    let previous = {
        let mut peer_state = peer_state_mutex.lock().await;
        let previous = peer_state.epoch.replace(remote);
        if previous.is_some_and(|previous| previous != remote) {
            // The peer switched documents; our sync state describes its old one. Count
            // the restart as sync activity, so the keepalive heal doesn't re-handshake
            // (resetting both sides again) while this fresh sync is under way.
            peer_state.reset_sync();
            peer_state.last_sync_rx = Some(Instant::now());
        }
        previous
    };

    // Log transitions only, not every keepalive.
    if previous != Some(remote) {
        match relation {
            EpochRelation::Same => {
                if previous.is_some() {
                    info!(peer = %peer, epoch = %remote, "Peer is on our epoch again; resuming sync");
                    let _ = state
                        .sync_events
                        .send(SyncEvent::PeerNeedsSync(peer.to_owned()));
                }
            }
            // Expected after a compaction: either side switches or catches up.
            EpochRelation::PeerAhead | EpochRelation::PeerBehind => {
                info!(
                    peer = %peer,
                    local_epoch = %local,
                    peer_epoch = %remote,
                    relation = ?relation,
                    "Peer is on another epoch; not syncing with it until both are on the same"
                );
            }
            EpochRelation::Diverged => {
                warn!(
                    peer = %peer,
                    local_epoch = %local,
                    peer_epoch = %remote,
                    "Peer is on a different epoch of the same generation; not syncing with it"
                );
            }
        }
    }

    relation == EpochRelation::Same
}

// --- inbound dispatch --------------------------------------------------------

pub(crate) async fn handle_message(
    state: &AppState,
    peer: String,
    message: transport::ParsedMessage,
) -> Result<()> {
    if peer == state.local_id() {
        debug!("Ignoring message from self");
        return Ok(());
    }

    remember_peer(state, &peer).await;

    // Any inbound message proves liveness.
    let reactivated = state
        .get_or_create_peer_state(&peer)
        .await
        .lock()
        .await
        .mark_seen(Instant::now());

    if reactivated {
        info!(peer = %peer, "Dormant peer is active again; re-handshaking");
    }

    let same_epoch = observe_peer_epoch(state, &peer, message.epoch()).await;

    //debug!(from = %peer, message = ?message, "Received message");

    match message {
        transport::ParsedMessage::Hello(_) => {
            // A HELLO means the peer is (re)starting its handshake with us, e.g. after a
            // restart. Whatever our sync state believed about it is stale.
            {
                let peer_state_mutex = state.get_or_create_peer_state(&peer).await;
                let mut peer_state = peer_state_mutex.lock().await;
                peer_state.reset_sync();
                peer_state.reply_owed = true;
            }

            // Sends now if we know our cluster ID; otherwise stays owed.
            send_owed_hello_reply(state, &peer).await;

            // Our own HELLO goes out regardless of whether the reply could be sent.
            if let Err(error) = send_hello_once(state, &peer).await {
                warn!(peer = %peer, error = %format!("{error:#}"), "HELLO failed; will retry");
            }

            // Restart sync with the fresh state (no-op until the peer is a validated member).
            let _ = state
                .sync_events
                .send(SyncEvent::PeerNeedsSync(peer.clone()));
        }
        transport::ParsedMessage::HelloReply(reply) => {
            let pending = establish_peer(state, &peer, reply.cluster_id.clone()).await?;

            discover_peers(state, &peer, &reply.peers).await;

            // Our cluster ID may have just become known: pay any replies we owe.
            flush_owed_hello_replies(state).await;

            if !pending.is_empty() {
                debug!(
                    to = %peer,
                    pending_messages = pending.len(),
                    "Replaying pending sync messages"
                );
                for envelope in pending {
                    if let Err(error) = handle_sync(state, &peer, envelope).await {
                        warn!(peer = %peer, error = %format!("{error:#}"), "Discarding invalid pending sync envelope");
                    }
                }
            }

            state
                .sync_events
                .send(SyncEvent::PeerValidated(peer.clone()))
                .map_err(|error| anyhow!("Failed to notify sync worker: {error}"))?;
        }
        transport::ParsedMessage::Sync(envelope) => {
            if reactivated {
                ensure_handshake(state, &peer).await;
            }
            if !same_epoch {
                debug!(peer = %peer, epoch = %envelope.epoch, "Dropping sync message from another epoch");
                return Ok(());
            }
            handle_sync(state, &peer, envelope).await?;
        }
        transport::ParsedMessage::Ping(remote) => {
            if reactivated || peer_cluster(state, &peer).await.is_none() {
                ensure_handshake(state, &peer).await;
            }
            let local = local_keepalive(state).await;
            heal_if_diverged(state, &peer, &local, &remote).await;
            if let Err(error) = state
                .transport
                .send_message(&peer, &ParsedMessage::Pong(local))
                .await
            {
                debug!(peer = %peer, error = %format!("{error:#}"), "PONG failed");
            }
        }
        transport::ParsedMessage::Pong(remote) => {
            if reactivated || peer_cluster(state, &peer).await.is_none() {
                ensure_handshake(state, &peer).await;
            }
            let local = local_keepalive(state).await;
            heal_if_diverged(state, &peer, &local, &remote).await;
        }
    }

    Ok(())
}