//! 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;

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);
        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)
    }

    /// 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 superseded(&self) -> crate::remake::Superseded {
        self.superseded_writes
    }

    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);
                    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 {
            tokio::task::spawn_blocking(move || {
                let result = {
                    let mut replica = state.doc.blocking_write();
                    let result = 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
        }
    }
}