//! Per-node, in-memory compaction state: snapshots built ahead of the switch, switch
//! progress for reports, a rejoin in progress, quarantined changes, and archived
//! documents of previous epochs.
//!
//! Nothing here is replicated or durable. After a restart the node simply has no
//! archive or quarantine, and rebuilds snapshots when needed.

use std::collections::HashMap;
use std::sync::{Mutex, MutexGuard};

use automerge::AutoCommit;
use serde::Serialize;
use tokio::time::Instant;

use super::records::{unix_millis, Retention};
use super::snapshot::Snapshot;
use crate::replica::Epoch;

#[derive(Default)]
pub(crate) struct Compaction {
    // A std mutex: every access is short and never spans an await.
    inner: Mutex<Inner>,
}

#[derive(Default)]
struct Inner {
    /// Built by the coordinator at commit time; consumed by the local switch.
    prebuilt: HashMap<String, Snapshot>,
    /// When this node first saw a compaction committed.
    first_seen: HashMap<String, Instant>,
    status: HashMap<String, LocalStatus>,
    archive: Vec<Archived>,
    /// Last attempt to catch up with a newer epoch.
    catch_up: Option<CatchUp>,
    /// A rejoin in progress: the node syncs an empty document of a newer epoch.
    rejoin: Option<Rejoin>,
    quarantine: Vec<Quarantined>,
}

/// Catching up with a newer epoch: attempts and the last failure.
#[derive(Clone, Debug)]
pub(crate) struct CatchUp {
    pub(crate) target: Epoch,
    pub(crate) attempts: u32,
    pub(crate) last_attempt: Instant,
    pub(crate) error: Option<String>,
}

/// The node left `previous` and is syncing `target` from scratch.
pub(crate) struct Rejoin {
    pub(crate) previous: Epoch,
    pub(crate) target: Epoch,
    /// The document of `previous`, until it is replayed, quarantined or restored.
    pub(crate) old: AutoCommit,
    pub(crate) started: Instant,
    /// The node served requests before; it does again once the rejoin is over.
    pub(crate) was_ready: bool,
}

/// Changes of this node that were not replayed into the current epoch.
pub(crate) struct Quarantined {
    pub(crate) id: String,
    /// The epoch the changes were made in, and the one the node rejoined.
    pub(crate) epoch: Epoch,
    pub(crate) target: Epoch,
    pub(crate) created_at: u64,
    pub(crate) reason: String,
    /// Store paths the changes touched (minimal, prefix-free).
    pub(crate) paths: Vec<Vec<String>>,
    /// The old document; `paths` are read from its merged view.
    pub(crate) document: AutoCommit,
}

#[derive(Clone, Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct CatchUpReport {
    pub(crate) target: String,
    pub(crate) attempts: u32,
    /// True while the node syncs the target epoch (not serving requests).
    pub(crate) rejoining: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub(crate) error: Option<String>,
}

/// What happened to a compaction on this node.
#[derive(Clone, Debug, Default, Serialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct LocalStatus {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub(crate) switched_at: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub(crate) replayed_paths: Option<usize>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub(crate) failed_paths: Option<usize>,
    /// ID of the quarantine entry holding this node's changes that weren't replayed.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub(crate) quarantine: Option<String>,
    /// Why the switch did not happen (yet).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub(crate) error: Option<String>,
}

struct Archived {
    epoch: Epoch,
    // Kept for auditing and recovery by hand; not read by the protocol.
    #[allow(dead_code)]
    doc: AutoCommit,
    archived_at: u64,
    retention: Retention,
}

/// Summary of an archived epoch.
#[derive(Clone, Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct ArchivedEpoch {
    pub(crate) generation: u64,
    pub(crate) id: String,
    pub(crate) archived_at: u64,
}

impl Compaction {
    fn lock(&self) -> MutexGuard<'_, Inner> {
        self.inner
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
    }

    pub(crate) fn store_prebuilt(&self, id: &str, snapshot: Snapshot) {
        self.lock().prebuilt.insert(id.to_owned(), snapshot);
    }

    pub(crate) fn take_prebuilt(&self, id: &str) -> Option<Snapshot> {
        self.lock().prebuilt.remove(id)
    }

    /// When this node first saw the compaction committed (now, on the first call).
    pub(crate) fn first_seen(&self, id: &str) -> Instant {
        *self
            .lock()
            .first_seen
            .entry(id.to_owned())
            .or_insert_with(Instant::now)
    }

    pub(crate) fn status(&self, id: &str) -> Option<LocalStatus> {
        self.lock().status.get(id).cloned()
    }

    pub(crate) fn set_status(&self, id: &str, status: LocalStatus) {
        self.lock().status.insert(id.to_owned(), status);
    }

    /// Keeps the previous epoch's document, then drops what has outlived its retention.
    pub(crate) fn archive(
        &self,
        epoch: Epoch,
        doc: AutoCommit,
        retention: Retention,
        current_generation: u64,
    ) {
        let mut inner = self.lock();
        inner.archive.push(Archived {
            epoch,
            doc,
            archived_at: unix_millis(),
            retention,
        });
        prune(&mut inner.archive, current_generation);
    }

    pub(crate) fn prune_archive(&self, current_generation: u64) {
        prune(&mut self.lock().archive, current_generation);
    }

    pub(crate) fn catch_up(&self) -> Option<CatchUp> {
        self.lock().catch_up.clone()
    }

    pub(crate) fn set_catch_up(&self, catch_up: Option<CatchUp>) {
        self.lock().catch_up = catch_up;
    }

    pub(crate) fn catch_up_report(&self) -> Option<CatchUpReport> {
        let inner = self.lock();
        let rejoining = inner.rejoin.is_some();
        inner.catch_up.as_ref().map(|catch_up| CatchUpReport {
            target: catch_up.target.to_string(),
            attempts: catch_up.attempts,
            rejoining,
            error: catch_up.error.clone(),
        })
    }

    pub(crate) fn is_rejoining(&self) -> bool {
        self.lock().rejoin.is_some()
    }

    /// Progress of the rejoin in progress: (target, started).
    pub(crate) fn rejoin_progress(&self) -> Option<(Epoch, Instant)> {
        self.lock()
            .rejoin
            .as_ref()
            .map(|rejoin| (rejoin.target, rejoin.started))
    }

    pub(crate) fn start_rejoin(&self, rejoin: Rejoin) {
        self.lock().rejoin = Some(rejoin);
    }

    pub(crate) fn take_rejoin(&self) -> Option<Rejoin> {
        self.lock().rejoin.take()
    }

    pub(crate) fn add_quarantine(&self, quarantined: Quarantined) {
        self.lock().quarantine.push(quarantined);
    }

    /// Runs `f` on the quarantined entry with this ID.
    pub(crate) fn with_quarantined<T>(
        &self,
        id: &str,
        f: impl FnOnce(&Quarantined) -> T,
    ) -> Option<T> {
        self.lock()
            .quarantine
            .iter()
            .find(|entry| entry.id == id)
            .map(f)
    }

    pub(crate) fn list_quarantine<T>(&self, f: impl Fn(&Quarantined) -> T) -> Vec<T> {
        self.lock().quarantine.iter().map(f).collect()
    }

    pub(crate) fn take_quarantined(&self, id: &str) -> Option<Quarantined> {
        let mut inner = self.lock();
        let index = inner.quarantine.iter().position(|entry| entry.id == id)?;
        Some(inner.quarantine.remove(index))
    }

    /// Puts an entry back, e.g. after a failed replay.
    pub(crate) fn restore_quarantined(&self, quarantined: Quarantined) {
        self.lock().quarantine.push(quarantined);
    }

    pub(crate) fn archived_epochs(&self) -> Vec<ArchivedEpoch> {
        self.lock()
            .archive
            .iter()
            .map(|archived| ArchivedEpoch {
                generation: archived.epoch.generation,
                id: archived.epoch.id_hex(),
                archived_at: archived.archived_at,
            })
            .collect()
    }
}

/// An archive is kept until every limit of its retention has passed.
fn prune(archive: &mut Vec<Archived>, current_generation: u64) {
    let now = unix_millis();
    archive.retain(|archived| {
        archived.retention.covers(
            current_generation.saturating_sub(archived.epoch.generation),
            now.saturating_sub(archived.archived_at),
        )
    });
}

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

    #[test]
    fn archive_retention() {
        let compaction = Compaction::default();
        let epoch = |generation| Epoch {
            generation,
            id: [generation as u8; 32],
        };

        // Kept by epoch count only.
        compaction.archive(
            epoch(0),
            AutoCommit::new(),
            Retention {
                epochs: Some(1),
                seconds: None,
            },
            1,
        );
        assert_eq!(compaction.archived_epochs().len(), 1);
        compaction.prune_archive(2);
        assert!(compaction.archived_epochs().is_empty());

        // Kept by time while the epoch count has passed.
        compaction.archive(
            epoch(1),
            AutoCommit::new(),
            Retention {
                epochs: Some(0),
                seconds: Some(3600),
            },
            5,
        );
        assert_eq!(compaction.archived_epochs().len(), 1);

        // No limits at all: not kept.
        compaction.archive(
            epoch(2),
            AutoCommit::new(),
            Retention {
                epochs: None,
                seconds: None,
            },
            3,
        );
        assert_eq!(compaction.archived_epochs().len(), 1);
    }
}