//! The node as the remake module sees it (`remake::Host`), and what follows an epoch
//! switch.

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

use anyhow::Result;
use tokio::sync::broadcast;
use tokio::time::Instant;
use tracing::{debug, error, info};

use super::{AppState, HelloState, NodeLifecycle, StaleView, Standing};
use crate::document::replica::{Epoch, Replica};
use crate::document::store::{ChangeSet, StoreError};
use crate::document::{ChangeHash, Document, ReadDoc, Transaction};
use crate::locks::RwLockExt;
use crate::remake::{BegunRejoin, Host, PeerSummary, RemakeState};

impl AppState {
    /// Under the document's write lock: `doc` of `epoch` replaces the live document,
    /// which is returned. Every peer's sync state is reset before the lock is released,
    /// so no message is generated from or applied to a state that describes the old
    /// document. (Lock order doc -> peer, as in sync generation.)
    fn swap_document(&self, replica: &mut Replica, epoch: Epoch, doc: Document) -> Document {
        let old = std::mem::replace(&mut replica.doc, doc);
        // The old document's signatures, of the changes the new one holds too (it may
        // be the old one without some, or another of its epoch); this node's replayed
        // into it, signed at once.
        replica.signatures = replica.signatures.kept_for(&replica.doc);
        if let Some(key) = &self.node_key {
            let replica = &mut *replica;
            replica.signatures.sign_own_latest(&replica.doc, key);
        }
        replica.epoch = epoch;
        replica.epoch_seq += 1;
        // Publishes the epoch with the heads (`Replica::stamps`), under the same lock,
        // so it never disagrees with the document for anyone holding the lock.
        replica.refresh_heads();
        let now = Instant::now();
        for peer in self.peer_states.blocking_read().values() {
            let mut peer = peer.blocking_lock();
            peer.reset_sync();
            // A fresh start, not an idle sync: no heal re-handshake for it.
            peer.last_sync_rx = Some(now);
        }
        old
    }

    /// Stops serving requests (back to `Joining`, keeping the swarm ID), as while a
    /// rejoin syncs a fresh document; on the blocking pool. Whether the node served.
    fn suspend(&self) -> bool {
        let mut lifecycle = self.lifecycle.blocking_write();
        let NodeLifecycle::Ready { swarm_id } = &*lifecycle else {
            return false;
        };
        let swarm_id = swarm_id.clone();
        debug!(swarm_id = %swarm_id, "Node suspended (rejoining)");
        *lifecycle = NodeLifecycle::Joining {
            discovered_swarm_id: Some(swarm_id),
        };
        true
    }

    /// Serves requests again after [`suspend`](Self::suspend); on the blocking pool.
    fn resume(&self) {
        let mut lifecycle = self.lifecycle.blocking_write();
        if let NodeLifecycle::Joining {
            discovered_swarm_id: Some(swarm_id),
        } = &*lifecycle
        {
            let swarm_id = swarm_id.clone();
            debug!(swarm_id = %swarm_id, "Node resumed");
            *lifecycle = NodeLifecycle::Ready { swarm_id };
        }
    }

    /// The document before a rejoin, out of the stale view (which ends) if the node
    /// keeps it there, else `old`; and whether it was kept there (the node was ready).
    fn take_stale(&self, old: Option<Document>) -> (Document, bool) {
        let Some(view) = self.stale.write_or_recover().take() else {
            return (old.unwrap_or_default(), false);
        };
        // A read holding it a moment longer gets a copy of its own.
        let doc = Arc::try_unwrap(view.doc)
            .map(|doc| {
                doc.into_inner()
                    .unwrap_or_else(|poisoned| poisoned.into_inner())
            })
            .unwrap_or_else(|shared| shared.read_or_recover().clone());
        (doc, true)
    }

    /// Drops `root`, a change the swarm's rules refuse that this node took unjudged
    /// (`swarm::refusals`), and every change built on it: the document becomes the one
    /// without them (forked where none is), written by a fresh actor of this node's, as
    /// the old one's later changes are among those dropped; this node's own writes among
    /// them go to the quarantine, the others' to the refused (`GET /admin/refused`, saying
    /// so). Every peer syncs anew. Nothing while a rejoin runs, or if the document doesn't
    /// hold `root`; nor if `root` is an epoch's root, or nothing would be kept, which no
    /// rule can mean (logged as an error). A drop taking much (`BIG_DROP`) waits for a
    /// second member to say so (`peer` the first, `why` its rules' reason): its rules and
    /// one peer's word might both be off, the swarm's not.
    pub(crate) async fn drop_taken(&self, root: ChangeHash, peer: &str, why: &str) {
        /// Changes a drop takes, at most, on one peer's word: more, or more than a tenth
        /// of the document (past `SMALL_DROP`), wait for another's.
        const BIG_DROP: usize = 100;
        const SMALL_DROP: usize = 10;
        let own = self.remake.own_actors();
        let state = self.clone();
        let peer = peer.to_owned();
        let why = why.to_owned();
        let dropped = tokio::task::spawn_blocking(move || {
            let mut replica = state.doc.blocking_write();
            if !replica.writable {
                return None;
            }
            let changes = replica.doc.get_changes(&[]);
            let dropped = changes.iter().find(|change| change.hash() == root)?;
            if dropped.deps().is_empty() {
                error!(change = %root, "Asked to drop the epoch's root; kept it");
                return None;
            }
            // In causal order: a change after those it builds on.
            let mut tainted: std::collections::HashSet<ChangeHash> = [root].into();
            for change in &changes {
                if change.deps().iter().any(|dep| tainted.contains(dep)) {
                    tainted.insert(change.hash());
                }
            }
            let kept: Vec<_> = changes
                .iter()
                .filter(|change| !tainted.contains(&change.hash()))
                .collect();
            if kept.is_empty() {
                error!(change = %root, "Dropping a change would leave nothing; kept it");
                return None;
            }
            if tainted.len() > BIG_DROP
                || (tainted.len() > SMALL_DROP && tainted.len() * 10 > changes.len())
            {
                use crate::locks::MutexExt;
                let mut screening = state.screening.lock_or_recover();
                let said = screening
                    .of(replica.epoch_seq)
                    .hold_drop(root, &peer);
                if said < 2 {
                    error!(change = %root, peer = %peer, changes = tainted.len(), why = %why, "A peer says we hold a change the swarm's rules refuse, and ours do too, but dropping it takes much: waiting for another member to say so");
                    return None;
                }
            }
            let built_on: std::collections::HashSet<ChangeHash> = kept
                .iter()
                .flat_map(|change| change.deps().iter().copied())
                .collect();
            let heads: Vec<ChangeHash> = kept
                .iter()
                .map(|change| change.hash())
                .filter(|hash| !built_on.contains(hash))
                .collect();
            // This node's own writes among those dropped: the paths they touched.
            let mut actors = own;
            actors.push(replica.doc.get_actor().clone());
            let ours: Vec<crate::document::ExpandedChange> = changes
                .iter()
                .filter(|change| tainted.contains(&change.hash()))
                .filter(|change| actors.contains(change.actor_id()))
                .map(|change| change.decode())
                .collect();
            let mut paths: Vec<Vec<String>> =
                crate::document::inspect::Inspector::new(&replica.doc)
                    .ops(&ours)
                    .into_iter()
                    .flatten()
                    .map(|op| op.path)
                    .collect();
            paths.sort();
            paths.dedup();
            let actor = crate::document::ActorId::random();
            let fresh = match replica.doc.fork_at(&heads) {
                Ok(fresh) => fresh.with_actor(actor.clone()),
                Err(error) => {
                    error!(error = %error, "Couldn't drop a change taken unjudged");
                    return None;
                }
            };
            // The others' changes dropped, for operators (own: to the quarantine).
            let theirs: Vec<crate::access::screen::Refused> = changes
                .iter()
                .filter(|change| tainted.contains(&change.hash()))
                .filter(|change| !actors.contains(change.actor_id()))
                .map(|change| crate::access::screen::Refused {
                    change: (*change).clone(),
                    committer: replica
                        .doc
                        .get_author_for_actor(change.actor_id())
                        .and_then(|author| {
                            std::str::from_utf8(author.as_bytes())
                                .ok()
                                .map(str::to_owned)
                        }),
                    user: change
                        .message()
                        .and_then(crate::document::replica::user_of)
                        .map(str::to_owned),
                    paths: Default::default(),
                    why: if change.hash() == root {
                        format!("taken unjudged, dropped: {why}")
                    } else {
                        format!("taken unjudged, dropped: it builds on {root}")
                    },
                })
                .collect();
            state.remake.add_own_actor(actor);
            let epoch = replica.epoch;
            let old = state.swap_document(&mut replica, epoch, fresh);
            {
                use crate::locks::MutexExt;
                let mut screening = state.screening.lock_or_recover();
                screening.of(replica.epoch_seq).keep_dropped(theirs);
            }
            Some((epoch, old, paths, tainted.len()))
        })
        .await
        .ok()
        .flatten();
        let Some((epoch, old, paths, count)) = dropped else {
            return;
        };
        info!(
            changes = count,
            own_paths = paths.len(),
            "Dropped a change taken unjudged and what built on it"
        );
        crate::remake::quarantine_taken(self, epoch, &old, paths);
        self.after_epoch_switch().await;
    }

    /// After the document was replaced: every peer's sync state refers to the old
    /// document, so start over with all of them. Peers still on the old epoch are
    /// skipped by the sync worker until they switch too.
    async fn after_epoch_switch(&self) {
        let peers = self.peers_snapshot().await;

        // Sync states were reset under the document lock by the swap itself.
        self.publish_change();

        // Tell peers about the new epoch now instead of with the next keepalive (up to a
        // keepalive interval later). Ready peers already on it learn it from our sync
        // opening, which goes out right away; everyone else gets a STATUS, so peers that
        // never saw the commit start catching up (we don't sync with them, so nothing
        // else would tell them).
        let local = self.local_epoch();
        let mut announce = Vec::new();
        for (id, peer_state_mutex) in &peers {
            // Others learn it from their own neighbours.
            if !crate::swarm::topology::is_neighbour(self, id) {
                continue;
            }
            let mut peer_state = peer_state_mutex.lock().await;
            // What it said about syncing with us was about our old epoch.
            peer_state.syncs_with_us = None;
            let gets_our_opening =
                peer_state.syncs_on(local) && peer_state.reports_ready == Some(true);
            if !gets_our_opening {
                announce.push(id);
            }
        }

        for id in announce {
            if let Err(error) = crate::swarm::peer::send_status(self, id).await {
                debug!(peer = %id, error = %format!("{error:#}"), "Epoch announcement failed");
            }
        }

        for (id, _) in peers {
            self.request_sync(&id);
        }
    }
}

/// The node as seen by the remake module.
impl Host for AppState {
    fn remake(&self) -> &RemakeState {
        &self.remake
    }

    fn clock(&self) -> Arc<crate::clock::Clock> {
        Arc::clone(&self.clock)
    }

    fn node_id(&self) -> String {
        self.local_id().to_owned()
    }

    fn current_epoch(&self) -> Epoch {
        self.local_epoch()
    }

    fn dormant_after(&self) -> Duration {
        self.timing.dormant_after
    }

    async fn clock_skew(&self) -> crate::swarm::skew::Skew {
        crate::swarm::skew::measure(self).await
    }

    fn electorate(&self, doc: &impl ReadDoc, now: crate::clock::UnixMillis) -> Vec<String> {
        crate::swarm::members::electorate(doc, now.as_u64(), self.electorate_window)
    }

    fn fetch_change(
        &self,
        hash: ChangeHash,
    ) -> impl Future<Output = Result<crate::document::Change>> + Send {
        let state = self.clone();
        async move {
            let peers = state.active_peers_on_local_epoch().await;
            crate::swarm::blob::fetch_change(&state, hash, peers).await
        }
    }

    fn fetch_content(
        &self,
        hash: [u8; 32],
        size: u64,
    ) -> impl Future<Output = Result<Vec<u8>>> + Send {
        let state = self.clone();
        async move {
            let peers = state.active_peers_on_local_epoch().await;
            crate::swarm::blob::fetch_content(&state, hash, size, peers).await
        }
    }

    fn request_sync(&self) {
        crate::swarm::source::resync_everyone(self);
    }

    fn is_node_ready(&self) -> impl Future<Output = bool> + Send {
        let state = self.clone();
        async move { state.is_ready().await }
    }

    fn peer_summaries(&self) -> impl Future<Output = Vec<PeerSummary>> + Send {
        let state = self.clone();
        async move {
            let peers = state.peers_snapshot().await;

            let mut summaries = Vec::with_capacity(peers.len());
            for (id, peer) in peers {
                let peer = peer.lock().await;
                summaries.push(PeerSummary {
                    id,
                    active: !peer.liveness.is_dormant(),
                    member: matches!(peer.standing, Standing::Member { .. })
                        && peer.hello_state == HelloState::Established,
                    ready: peer.reports_ready,
                    epoch: peer.epoch,
                });
            }
            summaries
        }
    }

    fn initial_sync_done(&self) -> bool {
        self.initial_sync_done
            .load(std::sync::atomic::Ordering::Acquire)
    }

    fn reset_initial_sync(&self) {
        self.initial_sync_done
            .store(false, std::sync::atomic::Ordering::Release);
        // The source chosen, and the one avoided, were for the catch-up before.
        crate::swarm::source::reset(self);
    }

    fn subscribe_changes(&self) -> broadcast::Receiver<Arc<ChangeSet>> {
        self.change_events.subscribe()
    }

    fn read_doc<T, F>(&self, f: F) -> impl Future<Output = T> + Send
    where
        F: FnOnce(&Document) -> T + Send + 'static,
        T: Send + 'static,
    {
        let state = self.clone();
        async move { state.render(f, |value| value).await }
    }

    fn commit_doc<T, E, F>(&self, f: F) -> impl Future<Output = Result<T, E>> + Send
    where
        F: FnOnce(&mut Transaction<'_>) -> Result<T, E> + Send + 'static,
        T: Send + 'static,
        E: From<StoreError> + Send + 'static,
    {
        let state = self.clone();
        async move { state.mutate_doc(f).await }
    }

    fn commit_doc_for<T, E, F>(
        &self,
        user: Option<String>,
        f: F,
    ) -> impl Future<Output = Result<T, E>> + Send
    where
        F: FnOnce(&mut Transaction<'_>) -> Result<T, E> + Send + 'static,
        T: Send + 'static,
        E: From<StoreError> + Send + 'static,
    {
        let state = self.clone();
        async move {
            state
                .mutate_doc_for(user, crate::document::signatures::Sign::Protocol, f)
                .await
                .map(|(out, _)| out)
        }
    }

    fn superseded(&self) -> crate::remake::Superseded {
        self.superseded_writes
    }

    fn strict_keys(&self) -> &[crate::access::sshsig::PublicKey] {
        &self.strict_keys
    }

    fn follows_strict_remakes(&self) -> bool {
        self.follow_strict_remakes
    }

    fn begin_rejoin(
        &self,
        target: Epoch,
        actor: crate::document::ActorId,
    ) -> impl Future<Output = BegunRejoin> + Send {
        let state = self.clone();
        async move {
            let begun = {
                let state = state.clone();
                tokio::task::spawn_blocking(move || {
                    let mut replica = state.doc.blocking_write();
                    let previous = replica.epoch;
                    let fresh = Document::new().with_actor(actor);
                    let old = state.swap_document(&mut replica, target, fresh);
                    replica.writable = false;
                    let heads = old.get_heads();
                    // Under the document's lock: no read or write in between finds the
                    // node ready on the fresh document.
                    let was_ready = state.suspend();
                    let old = if was_ready {
                        state.set_stale(Some(StaleView {
                            epoch: previous,
                            doc: Arc::new(std::sync::RwLock::new(old)),
                        }));
                        None
                    } else {
                        Some(old)
                    };
                    BegunRejoin {
                        was_ready,
                        heads,
                        old,
                    }
                })
                .await
                .expect("rejoin task panicked")
            };
            state.after_epoch_switch().await;
            begun
        }
    }

    fn rejoin_document(&self) -> impl Future<Output = Option<Document>> + Send {
        let state = self.clone();
        async move {
            let (_, stale) = state.stale_document()?;
            tokio::task::spawn_blocking(move || stale.read_or_recover().clone())
                .await
                .ok()
        }
    }

    fn finish_rejoin<T, F>(
        &self,
        old: Option<Document>,
        f: F,
    ) -> impl Future<Output = (T, Document)> + Send
    where
        F: FnOnce(&mut Document, &Document) -> T + Send + 'static,
        T: Send + 'static,
    {
        let state = self.clone();
        async move {
            let (out, old, resumed) = {
                let state = state.clone();
                tokio::task::spawn_blocking(move || {
                    let mut replica = state.doc.blocking_write();
                    // No write goes into the document before the rejoin any more: they
                    // take the document's lock first (`mutate_doc`).
                    let (old, resumed) = state.take_stale(old);
                    let out = f(&mut replica.doc, &old);
                    // What it replayed, vouched for at once.
                    if let Some(key) = &state.node_key {
                        let replica = &mut *replica;
                        replica.signatures.sign_own_latest(&replica.doc, key);
                    }
                    replica.writable = true;
                    replica.refresh_heads();
                    if resumed {
                        state.resume();
                    }
                    (out, old, resumed)
                })
                .await
                .expect("rejoin task panicked")
            };
            state.publish_change();
            state.join_progress.notify_one();
            if resumed {
                crate::swarm::peer::announce(&state);
            }
            (out, old)
        }
    }

    fn abandon_rejoin(
        &self,
        previous: Epoch,
        old: Option<Document>,
    ) -> impl Future<Output = ()> + Send {
        let state = self.clone();
        async move {
            let resumed = {
                let state = state.clone();
                tokio::task::spawn_blocking(move || {
                    let mut replica = state.doc.blocking_write();
                    let (old, resumed) = state.take_stale(old);
                    state.swap_document(&mut replica, previous, old);
                    replica.writable = true;
                    if resumed {
                        state.resume();
                    }
                    resumed
                })
                .await
                .expect("rejoin task panicked")
            };
            state.after_epoch_switch().await;
            state.join_progress.notify_one();
            if resumed {
                crate::swarm::peer::announce(&state);
            }
        }
    }

    fn apply_change(
        &self,
        change: crate::document::Change,
    ) -> impl Future<Output = Result<(), String>> + Send {
        let state = self.clone();
        async move {
            // Downloaded: looked at before Automerge applies it (`misplaced_author`).
            if let Some(misplaced) = crate::document::changes::misplaced_author(change.raw_bytes())
            {
                return Err(misplaced.to_string());
            }
            tokio::task::spawn_blocking(move || {
                let result = {
                    let mut replica = state.doc.blocking_write();
                    let policy = crate::access::policy::of(&replica.doc);
                    let result = if policy.enforced() {
                        // Judged first (`access::screen`).
                        let (epoch, epoch_seq) = (replica.epoch, replica.epoch_seq);
                        let replica = &mut *replica;
                        crate::access::screen::apply_bytes(
                            &mut replica.doc,
                            epoch,
                            epoch_seq,
                            change.raw_bytes(),
                            &state.screening,
                            policy.required().then_some(&mut replica.signatures),
                            &mut crate::document::PatchLog::inactive(),
                        )
                        .map(|_| ())
                        .map_err(|error| format!("{error:#}"))
                    } else {
                        replica
                            .apply_changes([change])
                            .map_err(|error| error.to_string())
                    };
                    replica.refresh_heads();
                    result
                }; // lock released before notifying

                if result.is_ok() {
                    state.publish_change();
                }
                result
            })
            .await
            .expect("document write task panicked")
        }
    }

    fn switch_epoch<E, F>(
        &self,
        epoch: Epoch,
        f: F,
    ) -> impl Future<Output = Result<Document, E>> + Send
    where
        F: FnOnce(&Document) -> Result<Document, E> + Send + 'static,
        E: Send + 'static,
    {
        let state = self.clone();
        async move {
            let swapped = {
                let state = state.clone();
                tokio::task::spawn_blocking(move || -> Result<Document, E> {
                    let mut replica = state.doc.blocking_write();
                    let new_doc = f(&replica.doc)?;
                    Ok(state.swap_document(&mut replica, epoch, new_doc))
                })
                .await
                .expect("epoch switch task panicked")
            };

            if swapped.is_ok() {
                state.after_epoch_switch().await;
            }
            swapped
        }
    }
}