use crate::{
    marshal::core::Variant,
    simplex::types::{Finalization, Notarization},
    types::{Epoch, Height, Round, View},
};
use commonware_codec::{CodecShared, Read};
use commonware_cryptography::{Digestible, certificate::Scheme};
use commonware_macros::boxed;
use commonware_runtime::{
    BufferPooler, Clock, Handle, Metrics, Spawner, Storage, buffer::paged::CacheRef,
};
use commonware_storage::{
    archive::{self, Archive as _, Identifier, MultiArchive as _, prunable},
    metadata::{self, Metadata},
    translator::TwoCap,
};
use rand_core::Rng;
use std::{
    cmp::max,
    collections::BTreeMap,
    future::Future,
    num::{NonZero, NonZeroUsize},
    time::Duration,
};
use tracing::{debug, info};

// The key used to store the current epoch in the metadata store.
const CACHED_EPOCHS_KEY: u8 = 0;

/// Configuration parameters for prunable archives.
pub(crate) struct Config {
    pub partition_prefix: String,
    pub prunable_items_per_section: NonZero<u64>,
    pub replay_buffer: NonZeroUsize,
    pub key_write_buffer: NonZeroUsize,
    pub value_write_buffer: NonZeroUsize,
    pub key_page_cache: CacheRef,
}

/// Prunable archives for a single epoch.
#[allow(clippy::type_complexity)]
struct Cache<R, V, S>
where
    R: BufferPooler + Rng + Spawner + Metrics + Clock + Storage,
    V: Variant,
    S: Scheme,
{
    /// Verified blocks stored by view
    verified_blocks: prunable::Archive<TwoCap, R, <V::Block as Digestible>::Digest, V::StoredBlock>,
    /// Notarized blocks stored by view
    notarized_blocks:
        prunable::Archive<TwoCap, R, <V::Block as Digestible>::Digest, V::StoredBlock>,
    /// Certified blocks indexed by height and keyed by digest.
    certified_blocks:
        prunable::Archive<TwoCap, R, <V::Block as Digestible>::Digest, V::StoredBlock>,
    /// Notarizations stored by view
    notarizations: prunable::Archive<
        TwoCap,
        R,
        <V::Block as Digestible>::Digest,
        Notarization<S, V::Commitment>,
    >,
    /// Finalizations stored by view
    finalizations: prunable::Archive<
        TwoCap,
        R,
        <V::Block as Digestible>::Digest,
        Finalization<S, V::Commitment>,
    >,
}

impl<R, V, S> Cache<R, V, S>
where
    R: BufferPooler + Rng + Spawner + Metrics + Clock + Storage,
    V: Variant,
    S: Scheme,
{
    /// Prune view-indexed archives to the given view.
    async fn prune_by_view(mut self, min_view: View) -> Self {
        (
            self.verified_blocks,
            self.notarized_blocks,
            self.notarizations,
            self.finalizations,
        ) = futures::try_join!(
            self.verified_blocks.prune(min_view.get()),
            self.notarized_blocks.prune(min_view.get()),
            self.notarizations.prune(min_view.get()),
            self.finalizations.prune(min_view.get()),
        )
        .unwrap_or_else(|e| panic!("failed to prune archives: {e}"));
        debug!(min_view = %min_view, "pruned archives");
        self
    }

    /// Prune height-indexed archives to the given height.
    async fn prune_by_height(mut self, min_height: Height) -> Self {
        self.certified_blocks = self
            .certified_blocks
            .prune(min_height.get())
            .await
            .expect("failed to prune certified blocks");
        self
    }
}

/// Manages prunable caches and their metadata.
pub(crate) struct Manager<R, V, S>
where
    R: BufferPooler + Rng + Spawner + Metrics + Clock + Storage,
    V: Variant,
    S: Scheme,
{
    /// Context
    context: R,

    /// Configuration for underlying prunable archives
    cfg: Config,

    /// Codec configuration for block type
    block_codec_config: <V::ApplicationBlock as Read>::Cfg,

    /// Metadata store for recording which epochs may have data. The value is a tuple of the floor
    /// and ceiling, the minimum and maximum epochs (inclusive) that may have data.
    metadata: Metadata<R, u8, (Epoch, Epoch)>,

    /// A map from epoch to its cache
    caches: BTreeMap<Epoch, Cache<R, V, S>>,
}

impl<R, V, S> Manager<R, V, S>
where
    R: BufferPooler + Rng + Spawner + Metrics + Clock + Storage,
    V: Variant,
    S: Scheme,
{
    /// Initialize the cache manager and its metadata store.
    #[boxed]
    pub(crate) async fn init(
        context: R,
        cfg: Config,
        block_codec_config: <V::ApplicationBlock as Read>::Cfg,
    ) -> Self {
        // Initialize metadata
        let metadata = Metadata::init(
            context.child("metadata"),
            metadata::Config {
                partition: format!("{}-metadata", cfg.partition_prefix),
                codec_config: ((), ()),
            },
        )
        .await
        .expect("failed to initialize metadata");

        // We don't eagerly initialize any epoch caches here, they will be
        // initialized on demand, otherwise there could be coordination issues
        // around the scheme provider.
        Self {
            context,
            cfg,
            block_codec_config,
            metadata,
            caches: BTreeMap::new(),
        }
    }

    /// Load all persisted epoch caches so that `find_block` can discover
    /// blocks written before the last shutdown.
    pub(crate) async fn load_persisted_epochs(mut self) -> Self {
        let (floor, ceiling) = self.get_metadata();
        for e in floor.get()..=ceiling.get() {
            let epoch = Epoch::new(e);
            if !self.caches.contains_key(&epoch) {
                self.init_epoch(epoch).await;
            }
        }
        self
    }

    /// Retrieve the epoch range that may have data.
    fn get_metadata(&self) -> (Epoch, Epoch) {
        self.metadata
            .get(&CACHED_EPOCHS_KEY)
            .cloned()
            .unwrap_or((Epoch::zero(), Epoch::zero()))
    }

    /// Set the epoch range that may have data.
    async fn set_metadata(mut self, floor: Epoch, ceiling: Epoch) -> Self {
        self.metadata = self
            .metadata
            .put_sync(CACHED_EPOCHS_KEY, (floor, ceiling))
            .await
            .expect("failed to write metadata");
        self
    }

    /// Runs `op` on the cache for `epoch`, initializing the epoch if it doesn't exist and
    /// reinserting the cache `op` returns.
    ///
    /// If the epoch is less than the minimum cached epoch, then it has already been pruned,
    /// and this will return `None` without running `op`.
    async fn with_epoch<T, Fut>(
        mut self,
        epoch: Epoch,
        op: impl FnOnce(Cache<R, V, S>) -> Fut,
    ) -> (Self, Option<T>)
    where
        Fut: Future<Output = (Cache<R, V, S>, T)>,
    {
        let cache = if let Some(cache) = self.caches.remove(&epoch) {
            cache
        } else {
            // If the epoch is less than the epoch floor, then it has already been pruned
            let (floor, ceiling) = self.get_metadata();
            if epoch < floor {
                return (self, None);
            }

            // Update the metadata (metadata-first is safe; init is idempotent)
            if epoch > ceiling {
                self = self.set_metadata(floor, epoch).await;
            }

            // Initialize the epoch
            self.init_epoch(epoch).await;
            self.caches.remove(&epoch).expect("epoch just initialized")
        };
        let (cache, out) = op(cache).await;
        self.caches.insert(epoch, cache);
        (self, Some(out))
    }

    /// Helper to initialize the cache for a given epoch.
    #[boxed]
    async fn init_epoch(&mut self, epoch: Epoch) {
        let context = self.context.child("cache").with_attribute("epoch", epoch);
        let (verified_blocks, notarized_blocks, certified_blocks, notarizations, finalizations) = futures::join!(
            Self::init_archive(
                &context,
                &self.cfg,
                epoch,
                "verified",
                self.block_codec_config.clone()
            ),
            Self::init_archive(
                &context,
                &self.cfg,
                epoch,
                "notarized",
                self.block_codec_config.clone()
            ),
            Self::init_archive(
                &context,
                &self.cfg,
                epoch,
                "certified",
                self.block_codec_config.clone()
            ),
            Self::init_archive(
                &context,
                &self.cfg,
                epoch,
                "notarizations",
                S::certificate_codec_config_unbounded(),
            ),
            Self::init_archive(
                &context,
                &self.cfg,
                epoch,
                "finalizations",
                S::certificate_codec_config_unbounded(),
            ),
        );
        let existing = self.caches.insert(
            epoch,
            Cache {
                verified_blocks,
                notarized_blocks,
                certified_blocks,
                notarizations,
                finalizations,
            },
        );
        assert!(existing.is_none(), "cache already exists for epoch {epoch}");
    }

    /// Helper to initialize an archive.
    async fn init_archive<T: CodecShared>(
        ctx: &R,
        cfg: &Config,
        epoch: Epoch,
        name: &'static str,
        codec_config: T::Cfg,
    ) -> prunable::Archive<TwoCap, R, <V::Block as Digestible>::Digest, T> {
        let start = ctx.current();
        let archive_cfg = prunable::Config {
            translator: TwoCap,
            metadata_partition: format!("{}-cache-{epoch}-{name}-metadata", cfg.partition_prefix),
            key_partition: format!("{}-cache-{epoch}-{name}-key", cfg.partition_prefix),
            key_page_cache: cfg.key_page_cache.clone(),
            value_partition: format!("{}-cache-{epoch}-{name}-value", cfg.partition_prefix),
            items_per_section: cfg.prunable_items_per_section,
            compression: None,
            codec_config,
            replay_buffer: cfg.replay_buffer,
            key_write_buffer: cfg.key_write_buffer,
            value_write_buffer: cfg.value_write_buffer,
        };
        let archive = prunable::Archive::init(ctx.child(name), archive_cfg)
            .await
            .unwrap_or_else(|_| panic!("failed to initialize {name} archive"));
        info!(elapsed = ?ctx.current().duration_since(start).unwrap_or(Duration::ZERO), "restored {name} archive");
        archive
    }

    /// Add a verify-stage candidate block to the prunable archive and start syncing it.
    ///
    /// The archive name is historical: callers may start this durability work
    /// after structural validation and before the application verdict is known.
    /// Consensus must not treat presence in this cache as application validity.
    ///
    /// No certificate pins a round to a single candidate at this stage: an
    /// equivocating leader can land one block at this view (possibly before a
    /// crash) while consensus later verifies a different one. Candidates are
    /// stored with multi-put semantics so a same-view collision cannot silently
    /// drop the new block while the returned handle vouches only for the old
    /// one. A digest already stored at this view is not duplicated, and the
    /// covering handle reports the durability of its existing write.
    pub(crate) async fn put_verified(
        mut self,
        round: Round,
        digest: <V::Block as Digestible>::Digest,
        block: V::StoredBlock,
    ) -> (Self, Handle<()>) {
        let view = round.view().get();
        let handle;
        (self, handle) = self
            .with_epoch(round.epoch(), |mut cache| async move {
                // Deduplicate against this view only: the same digest may legitimately
                // be stored again at a later view (boundary re-proposal), and each view
                // needs its own copy to survive per-view retention pruning.
                let exists = match cache.verified_blocks.has_at(view, &digest).await {
                    Ok(exists) => exists,
                    Err(e) => panic!("failed to check verified blocks: {e}"),
                };
                let handle;
                if exists {
                    (cache.verified_blocks, handle) = Self::handle_start_result(
                        cache.verified_blocks.start_sync().await,
                        round,
                        "verified",
                    );
                } else {
                    let result = cache
                        .verified_blocks
                        .put_multi_start_sync(view, digest, block)
                        .await;
                    (cache.verified_blocks, handle) =
                        Self::handle_start_result(result, round, "verified");
                }
                (cache, handle)
            })
            .await;
        (self, handle.unwrap_or_else(|| Handle::ready(Ok(()))))
    }

    /// Add a certified block to the height-indexed archive.
    pub(crate) async fn put_certified(
        mut self,
        epoch: Epoch,
        height: Height,
        digest: <V::Block as Digestible>::Digest,
        block: V::StoredBlock,
    ) -> Self {
        (self, _) = self
            .with_epoch(epoch, |mut cache| async move {
                // A digest determines its height, so scoping the dedup to this height
                // is exact and avoids fetching values.
                let exists = match cache.certified_blocks.has_at(height.get(), &digest).await {
                    Ok(exists) => exists,
                    Err(e) => panic!("failed to check certified block: {e}"),
                };
                if !exists {
                    cache.certified_blocks = cache
                        .certified_blocks
                        .put_multi_sync(height.get(), digest, block)
                        .await
                        .unwrap_or_else(|e| panic!("failed to insert certified block: {e}"));
                    debug!(%height, "cached certified block");
                }
                (cache, ())
            })
            .await;
        self
    }

    /// Add a notarized block to the prunable archive and start syncing it.
    pub(crate) async fn put_notarized(
        mut self,
        round: Round,
        digest: <V::Block as Digestible>::Digest,
        block: V::StoredBlock,
    ) -> (Self, Handle<()>) {
        let view = round.view().get();
        let handle;
        (self, handle) = self
            .with_epoch(round.epoch(), |mut cache| async move {
                let result = cache
                    .notarized_blocks
                    .put_start_sync(view, digest, block)
                    .await;
                let handle;
                (cache.notarized_blocks, handle) =
                    Self::handle_start_result(result, round, "notarized");
                (cache, handle)
            })
            .await;
        (self, handle.unwrap_or_else(|| Handle::ready(Ok(()))))
    }

    /// Returns a handle covering every write accepted by the round's verified-block
    /// archive before this call, including writes whose sync is still in flight.
    ///
    /// An absent epoch has nothing to observe (it never accepted a write or was
    /// pruned below the epoch floor), so its handle resolves immediately.
    pub(crate) async fn start_sync_verified(mut self, round: Round) -> (Self, Handle<()>) {
        let epoch = round.epoch();
        let Some(mut cache) = self.caches.remove(&epoch) else {
            return (self, Handle::ready(Ok(())));
        };
        let handle;
        (cache.verified_blocks, handle) =
            Self::handle_start_result(cache.verified_blocks.start_sync().await, round, "verified");
        self.caches.insert(epoch, cache);
        (self, handle)
    }

    /// Returns a handle covering every write accepted by the round's notarization
    /// archive before this call, including writes whose sync is still in flight.
    ///
    /// An absent epoch has nothing to observe (it never accepted a write or was
    /// pruned below the epoch floor), so its handle resolves immediately.
    pub(crate) async fn start_sync_notarizations(mut self, round: Round) -> (Self, Handle<()>) {
        let epoch = round.epoch();
        let Some(mut cache) = self.caches.remove(&epoch) else {
            return (self, Handle::ready(Ok(())));
        };
        let handle;
        (cache.notarizations, handle) = Self::handle_start_result(
            cache.notarizations.start_sync().await,
            round,
            "notarization",
        );
        self.caches.insert(epoch, cache);
        (self, handle)
    }

    /// Add a notarization to the prunable archive and start syncing it.
    pub(crate) async fn put_notarization(
        mut self,
        round: Round,
        digest: <V::Block as Digestible>::Digest,
        notarization: Notarization<S, V::Commitment>,
    ) -> (Self, Handle<()>) {
        let view = round.view().get();
        let handle;
        (self, handle) = self
            .with_epoch(round.epoch(), |mut cache| async move {
                let result = cache
                    .notarizations
                    .put_start_sync(view, digest, notarization)
                    .await;
                let handle;
                (cache.notarizations, handle) =
                    Self::handle_start_result(result, round, "notarization");
                (cache, handle)
            })
            .await;
        (self, handle.unwrap_or_else(|| Handle::ready(Ok(()))))
    }

    /// Add a finalization to the prunable archive.
    ///
    /// The blocking sync is intentional. A downstream application may write to
    /// an external store once it observes a finalization's effects and then
    /// assume the certificate is still readable from marshal after a restart.
    /// Deferring this sync would silently break that recovery pattern.
    pub(crate) async fn put_finalization(
        mut self,
        round: Round,
        digest: <V::Block as Digestible>::Digest,
        finalization: Finalization<S, V::Commitment>,
    ) -> Self {
        let view = round.view().get();
        (self, _) = self
            .with_epoch(round.epoch(), |mut cache| async move {
                cache.finalizations = cache
                    .finalizations
                    .put_sync(view, digest, finalization)
                    .await
                    .unwrap_or_else(|e| panic!("failed to insert finalization: {e}"));
                debug!(?round, "cached finalization");
                (cache, ())
            })
            .await;
        self
    }

    /// Helper to debug cache sync start results.
    fn handle_start_result<A>(
        result: Result<(A, Handle<()>), archive::Error>,
        round: Round,
        name: &str,
    ) -> (A, Handle<()>) {
        match result {
            Ok((archive, handle)) => {
                debug!(?round, name, "cache sync started");
                (archive, handle)
            }
            Err(e) => {
                panic!("failed to persist {name}: {e}");
            }
        }
    }

    /// Get a notarization from the prunable archive by round.
    pub(crate) async fn get_notarization(
        &self,
        round: Round,
    ) -> Option<Notarization<S, V::Commitment>> {
        let cache = self.caches.get(&round.epoch())?;
        cache
            .notarizations
            .get(Identifier::Index(round.view().get()))
            .await
            .expect("failed to get notarization")
    }

    /// Returns whether the verified archive holds `digest` at `round`.
    pub(crate) async fn has_verified(
        &self,
        round: Round,
        digest: &<V::Block as Digestible>::Digest,
    ) -> bool {
        let Some(cache) = self.caches.get(&round.epoch()) else {
            return false;
        };
        cache
            .verified_blocks
            .has_at(round.view().get(), digest)
            .await
            .expect("failed to check verified blocks")
    }

    /// Get the block previously persisted in the verified archive for `round`.
    ///
    /// The archive can hold multiple candidates at one view (an equivocating
    /// leader can land one before a crash and another after), and this returns
    /// the first stored. Callers must not assume it is the most recently
    /// verified candidate: check context/digest before reuse, or look up by
    /// digest.
    pub(crate) async fn get_verified(&self, round: Round) -> Option<V::StoredBlock> {
        let cache = self.caches.get(&round.epoch())?;
        cache
            .verified_blocks
            .get(Identifier::Index(round.view().get()))
            .await
            .expect("failed to get verified block")
    }

    /// Get a finalization from the prunable archive by block digest.
    ///
    /// SAFETY: For blocks/certificates admitted by marshal verification, a block digest
    /// maps to exactly one consensus payload commitment for the active marshal
    /// [`Variant`] instance.
    pub(crate) async fn get_finalization_for(
        &self,
        digest: <V::Block as Digestible>::Digest,
    ) -> Option<Finalization<S, V::Commitment>> {
        for cache in self.caches.values().rev() {
            match cache.finalizations.get(Identifier::Key(&digest)).await {
                Ok(Some(finalization)) => return Some(finalization),
                Ok(None) => continue,
                Err(e) => panic!("failed to get cached finalization: {e}"),
            }
        }
        None
    }

    /// Looks for a block (certified by height, verified, or notarized) that matches `predicate`.
    pub(crate) async fn find_block_matching(
        &self,
        digest: <V::Block as Digestible>::Digest,
        mut predicate: impl FnMut(&V::StoredBlock) -> bool,
    ) -> Option<V::StoredBlock> {
        // Check in reverse order
        for cache in self.caches.values().rev() {
            // Check verified blocks
            if let Some(block) = cache
                .verified_blocks
                .get(Identifier::Key(&digest))
                .await
                .expect("failed to get verified block")
                && predicate(&block)
            {
                return Some(block);
            }

            // Check notarized blocks
            if let Some(block) = cache
                .notarized_blocks
                .get(Identifier::Key(&digest))
                .await
                .expect("failed to get notarized block")
                && predicate(&block)
            {
                return Some(block);
            }

            // Check certified blocks
            if let Some(block) = cache
                .certified_blocks
                .get(Identifier::Key(&digest))
                .await
                .expect("failed to get certified block")
                && predicate(&block)
            {
                return Some(block);
            }
        }
        None
    }

    /// Prune the view-indexed caches below the given round.
    pub(crate) async fn prune_by_view(mut self, round: Round) -> Self {
        // Remove and close prunable archives from older epochs
        let new_floor = round.epoch();
        let old_epochs: Vec<Epoch> = self
            .caches
            .keys()
            .copied()
            .filter(|epoch| *epoch < new_floor)
            .collect();
        for epoch in old_epochs.iter() {
            let Cache {
                verified_blocks: vb,
                notarized_blocks: nb,
                certified_blocks: cb,
                notarizations: nv,
                finalizations: fv,
            } = self.caches.remove(epoch).unwrap();
            vb.destroy().await.expect("failed to destroy vb");
            nb.destroy().await.expect("failed to destroy nb");
            cb.destroy().await.expect("failed to destroy cb");
            nv.destroy().await.expect("failed to destroy nv");
            fv.destroy().await.expect("failed to destroy fv");
        }

        // Update metadata if necessary
        let (floor, ceiling) = self.get_metadata();
        if new_floor > floor {
            let new_ceiling = max(ceiling, new_floor);
            self = self.set_metadata(new_floor, new_ceiling).await;
        }

        // Prune archives for the given epoch
        let min_view = round.view();
        if let Some(prunable) = self.caches.remove(&round.epoch()) {
            let prunable = prunable.prune_by_view(min_view).await;
            self.caches.insert(round.epoch(), prunable);
        }
        self
    }

    /// Prune height-indexed certified blocks below the given height.
    pub(crate) async fn prune_by_height(mut self, height: Height) -> Self {
        for (epoch, cache) in std::mem::take(&mut self.caches) {
            let cache = cache.prune_by_height(height).await;
            self.caches.insert(epoch, cache);
        }
        self
    }
}