//! Serving reads: from the live document, or from the pre-rejoin copy while rejoining.

use std::sync::Arc;

use automerge::{AutoCommit, ChangeHash};

use super::AppState;
use crate::document::replica::Epoch;

/// A read-only view of the document the node served before it started rejoining.
pub(crate) struct StaleView {
    pub(crate) epoch: Epoch,
    pub(crate) doc: Arc<AutoCommit>,
}

/// Where a read was served from.
pub(crate) enum ReadSource {
    Live,
    /// The pre-rejoin document of this epoch.
    Stale(Epoch),
}

impl AppState {
    /// Renders a read for the store API: from the live document when the node is ready,
    /// from the stale view while it rejoins, and otherwise `503`, or the live document
    /// anyway with `unready_live`.
    ///
    /// With `with_heads`, a live read also returns the document's heads at the moment of
    /// the read (for clients that base later requests on what they read). Stale reads
    /// have no heads.
    pub(crate) async fn render_read<T, R>(
        &self,
        unready_live: bool,
        with_heads: bool,
        read: impl FnOnce(&AutoCommit) -> T + Send + 'static,
        finish: impl FnOnce(T) -> R + Send + 'static,
    ) -> Result<(R, ReadSource, Option<Vec<ChangeHash>>), axum::http::StatusCode>
    where
        T: Send + 'static,
        R: Send + 'static,
    {
        if self.is_ready().await {
            let (rendered, heads) = self.render_live(with_heads, read, finish).await;
            return Ok((rendered, ReadSource::Live, heads));
        }

        let stale = self
            .stale
            .read()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
            .as_ref()
            .map(|view| (view.epoch, Arc::clone(&view.doc)));
        let Some((epoch, doc)) = stale else {
            if unready_live {
                let (rendered, heads) = self.render_live(with_heads, read, finish).await;
                return Ok((rendered, ReadSource::Live, heads));
            }
            return Err(axum::http::StatusCode::SERVICE_UNAVAILABLE);
        };

        let rendered = tokio::task::spawn_blocking(move || finish(read(&doc)))
            .await
            .expect("document read task panicked");
        Ok((rendered, ReadSource::Stale(epoch), None))
    }

    async fn render_live<T, R>(
        &self,
        with_heads: bool,
        read: impl FnOnce(&AutoCommit) -> T + Send + 'static,
        finish: impl FnOnce(T) -> R + Send + 'static,
    ) -> (R, Option<Vec<ChangeHash>>)
    where
        T: Send + 'static,
        R: Send + 'static,
    {
        if !with_heads {
            return (self.render(read, finish).await, None);
        }
        // Heads and data under the same read lock (the heads are cached: `Replica`).
        let doc = Arc::clone(&self.doc);
        let (data, heads) = tokio::task::spawn_blocking(move || {
            let replica = doc.blocking_read();
            (read(&replica.doc), replica.heads().to_vec())
        })
        .await
        .expect("document read task panicked");
        let rendered = tokio::task::spawn_blocking(move || finish(data))
            .await
            .expect("document read task panicked");
        (rendered, Some(heads))
    }

    pub(super) fn set_stale(&self, view: Option<StaleView>) {
        *self
            .stale
            .write()
            .unwrap_or_else(|poisoned| poisoned.into_inner()) = view;
    }
}