//! 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 automerge::{AutoCommit, ChangeHash};
use tokio::sync::{Mutex, broadcast};
use tokio::time::Instant;
use tracing::debug;

use super::{AppState, HelloState, Liveness, NodeLifecycle, PeerState, StaleView, Standing};
use crate::document::replica::Epoch;
use crate::remake::{Host, PeerSummary, RemakeState};
use crate::swarm::sync::{ChangeSet, SyncEvent};

impl AppState {
    /// 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: Vec<(String, Arc<Mutex<PeerState>>)> = self
            .peer_states
            .read()
            .await
            .iter()
            .map(|(id, peer)| (id.clone(), Arc::clone(peer)))
            .collect();

        // 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 PING, 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.liveness == Liveness::Active
                && matches!(peer_state.standing, Standing::Member { .. })
                && peer_state.epoch == Some(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_ping(self, id).await {
                debug!(peer = %id, error = %format!("{error:#}"), "Epoch announcement failed");
            }
        }

        for (id, _) in peers {
            let _ = self.sync_events.send(SyncEvent::PeerNeedsSync(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
    }

    fn electorate(&self, doc: &AutoCommit, now_ms: u64) -> Vec<String> {
        crate::swarm::members::electorate(doc, now_ms, self.electorate_window)
    }

    fn fetch_change(
        &self,
        hash: ChangeHash,
    ) -> impl Future<Output = Result<automerge::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]) -> 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, 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: Vec<(String, Arc<Mutex<PeerState>>)> = state
                .peer_states
                .read()
                .await
                .iter()
                .map(|(id, peer)| (id.clone(), Arc::clone(peer)))
                .collect();

            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 == Liveness::Active,
                    member: matches!(peer.standing, Standing::Member { .. })
                        && peer.hello_state == HelloState::Established,
                    ready: peer.reports_ready,
                    epoch: peer.epoch,
                });
            }
            summaries
        }
    }

    fn suspend(&self) -> impl Future<Output = bool> + Send {
        let state = self.clone();
        async move {
            let mut lifecycle = state.lifecycle.write().await;
            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
        }
    }

    fn resume(&self) -> impl Future<Output = ()> + Send {
        let state = self.clone();
        async move {
            let mut lifecycle = state.lifecycle.write().await;
            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 };
            }
        }
    }

    fn set_stale_view(&self, view: Option<(Epoch, AutoCommit)>) {
        self.set_stale(view.map(|(epoch, doc)| StaleView {
            epoch,
            doc: Arc::new(doc),
        }));
    }

    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);
        *self
            .catch_up_source
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner()) = None;
        *self
            .catch_up_preference
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner()) = None;
    }

    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(&AutoCommit) -> 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 AutoCommit) -> Result<T, E> + Send + 'static,
        T: Send + 'static,
        E: Send + 'static,
    {
        let state = self.clone();
        async move { state.mutate_doc(f).await }
    }

    fn switch_epoch<E, F>(
        &self,
        epoch: Epoch,
        f: F,
    ) -> impl Future<Output = Result<AutoCommit, E>> + Send
    where
        F: FnOnce(&mut AutoCommit) -> Result<AutoCommit, E> + Send + 'static,
        E: Send + 'static,
    {
        let state = self.clone();
        async move {
            let swapped = {
                let state = state.clone();
                tokio::task::spawn_blocking(move || -> Result<AutoCommit, E> {
                    let mut replica = state.doc.blocking_write();
                    let new_doc = f(&mut replica.doc)?;
                    let old_doc = std::mem::replace(&mut replica.doc, new_doc);
                    replica.epoch = epoch;
                    replica.epoch_seq += 1;
                    replica.refresh_heads();
                    // Updated under the same lock, so the cached epoch never disagrees
                    // with the document for anyone holding the lock.
                    state.epoch_tx.send_replace(epoch);
                    // Reset every peer's sync state 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.)
                    let now = Instant::now();
                    for peer in state.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);
                    }
                    Ok(old_doc)
                })
                .await
                .expect("epoch switch task panicked")
            };

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