use super::round::Round;
use crate::{
    simplex::{
        elector::{Config as ElectorConfig, Elector},
        interesting, min_active,
        scheme::Scheme,
        types::{
            Artifact, Certificate, Context, Finalization, Finalize, Notarization, Notarize,
            Nullification, Nullify, Proposal,
        },
    },
    types::{Epoch, Participant, Round as Rnd, View, ViewDelta},
    Viewable,
};
use commonware_cryptography::{certificate, Digest};
use commonware_runtime::{telemetry::metrics::status::GaugeExt, Clock, Metrics};
use commonware_utils::futures::Aborter;
use prometheus_client::metrics::{counter::Counter, gauge::Gauge};
use rand_core::CryptoRngCore;
use std::{
    collections::{BTreeMap, BTreeSet},
    mem::{replace, take},
    sync::atomic::AtomicI64,
    time::{Duration, SystemTime},
};
use tracing::{debug, warn};

/// The view number of the genesis block.
const GENESIS_VIEW: View = View::zero();

/// Configuration for initializing [`State`].
pub struct Config<S: certificate::Scheme, L: ElectorConfig<S>> {
    pub scheme: S,
    pub elector: L,
    pub epoch: Epoch,
    pub activity_timeout: ViewDelta,
    pub leader_timeout: Duration,
    pub notarization_timeout: Duration,
    pub nullify_retry: Duration,
}

/// Per-[Epoch] state machine.
///
/// Tracks proposals and certificates for each view. Vote aggregation and verification
/// is handled by the [crate::simplex::actors::batcher].
pub struct State<E: Clock + CryptoRngCore + Metrics, S: Scheme<D>, L: ElectorConfig<S>, D: Digest> {
    context: E,
    scheme: S,
    elector: L::Elector,
    epoch: Epoch,
    activity_timeout: ViewDelta,
    leader_timeout: Duration,
    notarization_timeout: Duration,
    nullify_retry: Duration,
    view: View,
    last_finalized: View,
    genesis: Option<D>,
    views: BTreeMap<View, Round<S, D>>,

    certification_candidates: BTreeSet<View>,
    outstanding_certifications: BTreeSet<View>,

    current_view: Gauge,
    tracked_views: Gauge,
    skipped_views: Counter,
}

impl<E: Clock + CryptoRngCore + Metrics, S: Scheme<D>, L: ElectorConfig<S>, D: Digest>
    State<E, S, L, D>
{
    pub fn new(context: E, cfg: Config<S, L>) -> Self {
        let current_view = Gauge::<i64, AtomicI64>::default();
        let tracked_views = Gauge::<i64, AtomicI64>::default();
        let skipped_views = Counter::default();
        context.register("current_view", "current view", current_view.clone());
        context.register("tracked_views", "tracked views", tracked_views.clone());
        context.register("skipped_views", "skipped views", skipped_views.clone());

        // Build elector with participants
        let elector = cfg.elector.build(cfg.scheme.participants());

        Self {
            context,
            scheme: cfg.scheme,
            elector,
            epoch: cfg.epoch,
            activity_timeout: cfg.activity_timeout,
            leader_timeout: cfg.leader_timeout,
            notarization_timeout: cfg.notarization_timeout,
            nullify_retry: cfg.nullify_retry,
            view: GENESIS_VIEW,
            last_finalized: GENESIS_VIEW,
            genesis: None,
            views: BTreeMap::new(),
            certification_candidates: BTreeSet::new(),
            outstanding_certifications: BTreeSet::new(),
            current_view,
            tracked_views,
            skipped_views,
        }
    }

    /// Seeds the state machine with the genesis payload and advances into view 1.
    pub fn set_genesis(&mut self, genesis: D) {
        self.genesis = Some(genesis);
        self.enter_view(GENESIS_VIEW.next());
        self.set_leader(GENESIS_VIEW.next(), None);
    }

    /// Returns the epoch managed by this state machine.
    pub const fn epoch(&self) -> Epoch {
        self.epoch
    }

    /// Returns the view currently being driven.
    pub const fn current_view(&self) -> View {
        self.view
    }

    /// Returns the highest finalized view we have observed.
    pub const fn last_finalized(&self) -> View {
        self.last_finalized
    }

    /// Returns the lowest view that must remain in memory to satisfy the activity timeout.
    pub const fn min_active(&self) -> View {
        min_active(self.activity_timeout, self.last_finalized)
    }

    /// Returns whether `pending` is still relevant for progress, optionally allowing future views.
    pub fn is_interesting(&self, pending: View, allow_future: bool) -> bool {
        interesting(
            self.activity_timeout,
            self.last_finalized,
            self.view,
            pending,
            allow_future,
        )
    }

    /// Returns true when the local signer is the participant with index `idx`.
    pub fn is_me(&self, idx: Participant) -> bool {
        self.scheme.me().is_some_and(|me| me == idx)
    }

    /// Advances the view and updates the leader.
    ///
    /// If `seed` is `None`, this **must** be the first view after genesis (view 1).
    /// For all subsequent views, a seed derived from the previous view's certificate
    /// must be provided.
    fn enter_view(&mut self, view: View) -> bool {
        if view <= self.view {
            return false;
        }

        let now = self.context.current();
        let leader_deadline = now + self.leader_timeout;
        let advance_deadline = now + self.notarization_timeout;

        let round = self.create_round(view);
        round.set_deadlines(leader_deadline, advance_deadline);
        self.view = view;

        // Update metrics
        let _ = self.current_view.try_set(view.get());
        true
    }

    /// Sets the leader for the given view if it is not already set.
    fn set_leader(&mut self, view: View, certificate: Option<&S::Certificate>) {
        let leader = self.elector.elect(Rnd::new(self.epoch, view), certificate);
        let round = self.create_round(view);
        if round.leader().is_some() {
            return;
        }
        round.set_leader(leader);
    }

    /// Ensures a round exists for the given view.
    fn create_round(&mut self, view: View) -> &mut Round<S, D> {
        self.views.entry(view).or_insert_with(|| {
            Round::new(
                self.scheme.clone(),
                Rnd::new(self.epoch, view),
                self.context.current(),
            )
        })
    }

    /// Returns the deadline for the next timeout (leader, notarization, or retry).
    pub fn next_timeout_deadline(&mut self) -> SystemTime {
        let now = self.context.current();
        let nullify_retry = self.nullify_retry;
        let round = self.create_round(self.view);
        round.next_timeout_deadline(now, nullify_retry)
    }

    /// Handle a timeout event for the current view.
    /// Returns the nullify vote and optionally an entry certificate for the previous view
    /// (if this is a retry timeout and we can construct one).
    pub fn handle_timeout(&mut self) -> (bool, Option<Nullify<S>>, Option<Certificate<S, D>>) {
        let view = self.view;
        let Some(retry) = self.create_round(view).construct_nullify() else {
            return (false, None, None);
        };
        let nullify = Nullify::sign::<D>(&self.scheme, Rnd::new(self.epoch, view));

        // If was retry, we need to get entry certificates for the previous view
        let entry_view = view.previous().unwrap_or(GENESIS_VIEW);
        if !retry || entry_view == GENESIS_VIEW {
            return (retry, nullify, None);
        }

        // Get the certificate for the previous view. Prefer finalizations since they are the
        // strongest proof available. Prefer nullifications over notarizations because a
        // nullification overwrites an uncertified notarization (if we only heard of notarizations,
        // we may never exit a view with an uncertifiable notarization).
        #[allow(clippy::option_if_let_else)]
        let cert = if let Some(finalization) = self.finalization(entry_view).cloned() {
            Some(Certificate::Finalization(finalization))
        } else if let Some(nullification) = self.nullification(entry_view).cloned() {
            Some(Certificate::Nullification(nullification))
        } else if let Some(notarization) = self.notarization(entry_view).cloned() {
            Some(Certificate::Notarization(notarization))
        } else {
            warn!(%entry_view, "entry certificate not found during timeout");
            None
        };
        (retry, nullify, cert)
    }

    /// Inserts a notarization certificate and prepares the next view's leader.
    ///
    /// Does not advance into the next view until certification passes.
    /// Adds to certification candidates if successful.
    pub fn add_notarization(
        &mut self,
        notarization: Notarization<S, D>,
    ) -> (bool, Option<S::PublicKey>) {
        let view = notarization.view();
        // Do not advance to the next view until the certification passes
        self.set_leader(view.next(), Some(&notarization.certificate));
        let result = self.create_round(view).add_notarization(notarization);
        if result.0 && view > self.last_finalized {
            self.certification_candidates.insert(view);
        }
        result
    }

    /// Inserts a nullification certificate and advances into the next view.
    pub fn add_nullification(&mut self, nullification: Nullification<S>) -> bool {
        let view = nullification.view();
        self.enter_view(view.next());
        self.set_leader(view.next(), Some(&nullification.certificate));
        self.create_round(view).add_nullification(nullification)
    }

    /// Inserts a finalization certificate, updates the finalized height, and advances the view.
    pub fn add_finalization(
        &mut self,
        finalization: Finalization<S, D>,
    ) -> (bool, Option<S::PublicKey>) {
        let view = finalization.view();
        if view > self.last_finalized {
            self.last_finalized = view;

            // Prune certification candidates at or below finalized view
            self.certification_candidates.retain(|v| *v > view);

            // Abort outstanding certifications at or below finalized view
            let keep = self.outstanding_certifications.split_off(&view.next());
            for v in replace(&mut self.outstanding_certifications, keep) {
                if let Some(round) = self.views.get_mut(&v) {
                    round.abort_certify();
                }
            }
        }

        self.enter_view(view.next());
        self.set_leader(view.next(), Some(&finalization.certificate));
        self.create_round(view).add_finalization(finalization)
    }

    /// Construct a notarize vote for this view when we're ready to sign.
    pub fn construct_notarize(&mut self, view: View) -> Option<Notarize<S, D>> {
        let candidate = self
            .views
            .get_mut(&view)
            .and_then(|round| round.construct_notarize().cloned())?;

        // Signing can only fail if we are a verifier, so we don't need to worry about
        // unwinding our broadcast toggle.
        Notarize::sign(&self.scheme, candidate)
    }

    /// Construct a finalize vote if the round provides a candidate.
    pub fn construct_finalize(&mut self, view: View) -> Option<Finalize<S, D>> {
        let candidate = self
            .views
            .get_mut(&view)
            .and_then(|round| round.construct_finalize().cloned())?;

        // Signing can only fail if we are a verifier, so we don't need to worry about
        // unwinding our broadcast toggle.
        Finalize::sign(&self.scheme, candidate)
    }

    /// Construct a notarization certificate once the round has quorum.
    pub fn broadcast_notarization(&mut self, view: View) -> Option<Notarization<S, D>> {
        self.views
            .get_mut(&view)
            .and_then(|round| round.broadcast_notarization())
    }

    /// Return a notarization certificate, if one exists.
    pub fn notarization(&self, view: View) -> Option<&Notarization<S, D>> {
        self.views.get(&view).and_then(|round| round.notarization())
    }

    /// Return a nullification certificate, if one exists.
    pub fn nullification(&self, view: View) -> Option<&Nullification<S>> {
        self.views
            .get(&view)
            .and_then(|round| round.nullification())
    }

    /// Return a finalization certificate, if one exists.
    pub fn finalization(&self, view: View) -> Option<&Finalization<S, D>> {
        self.views.get(&view).and_then(|round| round.finalization())
    }

    /// Construct a nullification certificate once the round has quorum.
    pub fn broadcast_nullification(&mut self, view: View) -> Option<Nullification<S>> {
        self.views
            .get_mut(&view)
            .and_then(|round| round.broadcast_nullification())
    }

    /// Construct a finalization certificate once the round has quorum.
    pub fn broadcast_finalization(&mut self, view: View) -> Option<Finalization<S, D>> {
        self.views
            .get_mut(&view)
            .and_then(|round| round.broadcast_finalization())
    }

    /// Replays a journaled artifact into the appropriate round during recovery.
    pub fn replay(&mut self, artifact: &Artifact<S, D>) {
        self.create_round(artifact.view()).replay(artifact);
    }

    /// Returns the leader index for `view` if we already entered it.
    pub fn leader_index(&self, view: View) -> Option<Participant> {
        self.views
            .get(&view)
            .and_then(|round| round.leader().map(|leader| leader.idx))
    }

    /// Returns how long `view` has been live based on the clock samples stored by its round.
    pub fn elapsed_since_start(&self, view: View) -> Option<Duration> {
        let now = self.context.current();
        self.views
            .get(&view)
            .map(|round| round.elapsed_since_start(now))
    }

    /// Immediately expires `view`, forcing its timeouts to trigger on the next tick.
    pub fn expire_round(&mut self, view: View) {
        let now = self.context.current();
        self.create_round(view).set_deadlines(now, now);

        // Update metrics
        self.skipped_views.inc();
    }

    /// Attempt to propose a new block.
    pub fn try_propose(&mut self) -> Option<Context<D, S::PublicKey>> {
        // Perform fast checks before lookback
        let view = self.view;
        if view == GENESIS_VIEW {
            return None;
        }
        if !self
            .views
            .get_mut(&view)
            .expect("view must exist")
            .should_propose()
        {
            return None;
        }

        // Look for parent
        let parent = self.find_parent(view);
        let (parent_view, parent_payload) = match parent {
            Ok(parent) => parent,
            Err(missing) => {
                debug!(%view, %missing, "missing parent during proposal");
                return None;
            }
        };
        let leader = self
            .views
            .get_mut(&view)
            .expect("view must exist")
            .try_propose()?;
        Some(Context {
            round: Rnd::new(self.epoch, view),
            leader: leader.key,
            parent: (parent_view, parent_payload),
        })
    }

    /// Records a locally constructed proposal once the automaton finishes building it.
    pub fn proposed(&mut self, proposal: Proposal<D>) -> bool {
        self.views
            .get_mut(&proposal.view())
            .map(|round| round.proposed(proposal))
            .unwrap_or(false)
    }

    /// Sets a proposal received from the batcher (leader's first notarize vote).
    ///
    /// Returns true if the proposal should trigger verification, false otherwise.
    pub fn set_proposal(&mut self, view: View, proposal: Proposal<D>) -> bool {
        self.create_round(view).set_proposal(proposal)
    }

    /// Attempt to verify a proposed block.
    ///
    /// Unlike during proposal, we don't use a verification opportunity
    /// to backfill missing certificates (a malicious proposer could
    /// ask us to fetch junk).
    #[allow(clippy::type_complexity)]
    pub fn try_verify(&mut self) -> Option<(Context<D, S::PublicKey>, Proposal<D>)> {
        let view = self.view;
        let (leader, proposal) = self.views.get(&view)?.should_verify()?;
        let parent_payload = self.parent_payload(&proposal)?;
        if !self.views.get_mut(&view)?.try_verify() {
            return None;
        }
        let context = Context {
            round: proposal.round,
            leader: leader.key,
            parent: (proposal.parent, parent_payload),
        };
        Some((context, proposal))
    }

    /// Marks proposal verification as complete when the peer payload validates.
    pub fn verified(&mut self, view: View) -> bool {
        self.views
            .get_mut(&view)
            .map(|round| round.verified())
            .unwrap_or(false)
    }

    /// Store the abort handle for an in-flight certification request.
    pub fn set_certify_handle(&mut self, view: View, handle: Aborter) {
        let Some(round) = self.views.get_mut(&view) else {
            return;
        };
        round.set_certify_handle(handle);
        self.outstanding_certifications.insert(view);
    }

    /// Takes all certification candidates and returns proposals ready for certification.
    pub fn certify_candidates(&mut self) -> Vec<Proposal<D>> {
        let candidates = take(&mut self.certification_candidates);
        candidates
            .into_iter()
            .filter_map(|view| {
                if view <= self.last_finalized {
                    return None;
                }
                self.views.get_mut(&view)?.try_certify()
            })
            .collect()
    }

    /// Marks proposal certification as complete and returns the notarization.
    ///
    /// Returns `None` if the view was already pruned. Otherwise returns the notarization
    /// regardless of success/failure.
    pub fn certified(&mut self, view: View, is_success: bool) -> Option<Notarization<S, D>> {
        let round = self.views.get_mut(&view)?;
        round.certified(is_success);

        // Remove from outstanding since certification is complete
        self.outstanding_certifications.remove(&view);

        // Get notarization before advancing state
        let notarization = round
            .notarization()
            .cloned()
            .expect("notarization must exist for certified view");

        if is_success {
            self.enter_view(view.next());
        } else {
            self.expire_round(view);
        }

        Some(notarization)
    }

    /// Drops any views that fall below the activity horizon and returns them for logging.
    pub fn prune(&mut self) -> Vec<View> {
        let min = self.min_active();
        let kept = self.views.split_off(&min);
        let removed = replace(&mut self.views, kept).into_keys().collect();

        // Update metrics
        let _ = self.tracked_views.try_set(self.views.len());
        removed
    }

    /// Returns the payload of the proposal if it is certified (including finalized).
    fn is_certified(&self, view: View) -> Option<&D> {
        // Special case for genesis view
        if view == GENESIS_VIEW {
            return Some(self.genesis.as_ref().expect("genesis must be present"));
        }

        // Check for explicit certification
        let round = self.views.get(&view)?;
        if round.finalization().is_some() || round.is_certified() {
            return Some(&round.proposal().expect("proposal must exist").payload);
        }
        None
    }

    /// Returns true if the view is nullified.
    fn is_nullified(&self, view: View) -> bool {
        // Special case for genesis view (although it should also not be in the views map).
        if view == GENESIS_VIEW {
            return false;
        }

        let round = match self.views.get(&view) {
            Some(round) => round,
            None => return false,
        };
        round.nullification().is_some()
    }

    /// Returns true if certification for the view was aborted due to finalization.
    #[cfg(test)]
    pub fn is_certify_aborted(&self, view: View) -> bool {
        self.views
            .get(&view)
            .is_some_and(|round| round.is_certify_aborted())
    }

    /// Finds the parent payload for a given view by walking backwards through
    /// the chain, skipping nullified views until finding a certified payload.
    fn find_parent(&self, view: View) -> Result<(View, D), View> {
        // If the view is the genesis view, consider it to be its own parent.
        let mut cursor = view.previous().unwrap_or(GENESIS_VIEW);

        loop {
            // Return the first certified (including finalized) parent.
            if let Some(parent) = self.is_certified(cursor) {
                return Ok((cursor, *parent));
            }

            // If the view is also not nullified, there is a gap in certificates.
            if !self.is_nullified(cursor) {
                return Err(cursor);
            }

            cursor = cursor.previous().expect("cursor must not wrap");
        }
    }

    /// Returns the payload of the proposal's parent if:
    /// - It is less-than the proposal view.
    /// - It is greater-than-or-equal-to the last finalized view.
    /// - It is notarized or finalized.
    /// - There exist nullifications for all views between it and the proposal view.
    fn parent_payload(&self, proposal: &Proposal<D>) -> Option<D> {
        // Sanity check that the parent view is less than the proposal view.
        let (view, parent) = (proposal.view(), proposal.parent);
        if view <= parent {
            return None;
        }

        // Ignore any requests for outdated parent views.
        if parent < self.last_finalized {
            return None;
        }

        // Check that there are nullifications for all views between the parent and the proposal view.
        if !View::range(parent.next(), view).all(|v| self.is_nullified(v)) {
            return None;
        }

        // May return `None` if the parent view is not yet either:
        // - notarized and certified
        // - finalized
        self.is_certified(parent).copied()
    }

    /// Returns the certificate for the parent of the proposal at the given view.
    pub fn parent_certificate(&mut self, view: View) -> Option<Certificate<S, D>> {
        let parent = {
            let view = self.views.get(&view)?.proposal()?.parent;
            self.views.get(&view)?
        };

        if let Some(f) = parent.finalization().cloned() {
            return Some(Certificate::Finalization(f));
        }
        if let Some(n) = parent.notarization().cloned() {
            return Some(Certificate::Notarization(n));
        }
        None
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::simplex::{
        elector::RoundRobin,
        scheme::ed25519,
        types::{Finalization, Finalize, Notarization, Notarize, Nullification, Nullify, Proposal},
    };
    use commonware_cryptography::{certificate::mocks::Fixture, sha256::Digest as Sha256Digest};
    use commonware_parallel::Sequential;
    use commonware_runtime::{deterministic, Runner};
    use commonware_utils::futures::AbortablePool;
    use std::time::Duration;

    fn test_genesis() -> Sha256Digest {
        Sha256Digest::from([0u8; 32])
    }

    #[test]
    fn certificate_candidates_respect_force_flag() {
        let runtime = deterministic::Runner::default();
        runtime.start(|mut context| async move {
            let namespace = b"ns".to_vec();
            let Fixture {
                schemes, verifier, ..
            } = ed25519::fixture(&mut context, &namespace, 4);
            let mut state = State::new(
                context,
                Config {
                    scheme: verifier.clone(),
                    elector: <RoundRobin>::default(),
                    epoch: Epoch::new(11),
                    activity_timeout: ViewDelta::new(6),
                    leader_timeout: Duration::from_secs(1),
                    notarization_timeout: Duration::from_secs(2),
                    nullify_retry: Duration::from_secs(3),
                },
            );
            state.set_genesis(test_genesis());

            // Add notarization
            let notarize_view = View::new(3);
            let notarize_round = Rnd::new(Epoch::new(11), notarize_view);
            let notarize_proposal =
                Proposal::new(notarize_round, GENESIS_VIEW, Sha256Digest::from([50u8; 32]));
            let notarize_votes: Vec<_> = schemes
                .iter()
                .map(|scheme| Notarize::sign(scheme, notarize_proposal.clone()).unwrap())
                .collect();
            let notarization =
                Notarization::from_notarizes(&verifier, notarize_votes.iter(), &Sequential)
                    .expect("notarization");
            state.add_notarization(notarization);

            // Produce candidate once
            assert!(state.broadcast_notarization(notarize_view).is_some());
            assert!(state.broadcast_notarization(notarize_view).is_none());
            assert!(state.notarization(notarize_view).is_some());

            // Add nullification
            let nullify_view = View::new(4);
            let nullify_round = Rnd::new(Epoch::new(11), nullify_view);
            let nullify_votes: Vec<_> = schemes
                .iter()
                .map(|scheme| {
                    Nullify::sign::<Sha256Digest>(scheme, nullify_round).expect("nullify")
                })
                .collect();
            let nullification =
                Nullification::from_nullifies(&verifier, &nullify_votes, &Sequential)
                    .expect("nullification");
            state.add_nullification(nullification);

            // Produce candidate once
            assert!(state.broadcast_nullification(nullify_view).is_some());
            assert!(state.broadcast_nullification(nullify_view).is_none());
            assert!(state.nullification(nullify_view).is_some());

            // Add finalization
            let finalize_view = View::new(5);
            let finalize_round = Rnd::new(Epoch::new(11), finalize_view);
            let finalize_proposal =
                Proposal::new(finalize_round, GENESIS_VIEW, Sha256Digest::from([51u8; 32]));
            let finalize_votes: Vec<_> = schemes
                .iter()
                .map(|scheme| Finalize::sign(scheme, finalize_proposal.clone()).unwrap())
                .collect();
            let finalization =
                Finalization::from_finalizes(&verifier, finalize_votes.iter(), &Sequential)
                    .expect("finalization");
            state.add_finalization(finalization);

            // Produce candidate once
            assert!(state.broadcast_finalization(finalize_view).is_some());
            assert!(state.broadcast_finalization(finalize_view).is_none());
            assert!(state.finalization(finalize_view).is_some());
        });
    }

    #[test]
    fn timeout_helpers_reuse_and_reset_deadlines() {
        let runtime = deterministic::Runner::default();
        runtime.start(|mut context| async move {
            let namespace = b"ns".to_vec();
            let Fixture { schemes, .. } = ed25519::fixture(&mut context, &namespace, 4);
            let local_scheme = schemes[0].clone(); // leader of view 1
            let retry = Duration::from_secs(3);
            let cfg = Config {
                scheme: local_scheme.clone(),
                elector: <RoundRobin>::default(),
                epoch: Epoch::new(4),
                activity_timeout: ViewDelta::new(2),
                leader_timeout: Duration::from_secs(1),
                notarization_timeout: Duration::from_secs(2),
                nullify_retry: retry,
            };
            let mut state = State::new(context.clone(), cfg);
            state.set_genesis(test_genesis());

            // Should return same deadline until something done
            let first = state.next_timeout_deadline();
            let second = state.next_timeout_deadline();
            assert_eq!(first, second, "cached deadline should be reused");

            // Handle timeout should return false (not a retry)
            let (was_retry, _, _) = state.handle_timeout();
            assert!(!was_retry, "first timeout is not a retry");

            // Set retry deadline
            context.sleep(Duration::from_secs(2)).await;
            let later = context.current();

            // Confirm retry deadline is set
            let third = state.next_timeout_deadline();
            assert_eq!(third, later + retry, "new retry scheduled after timeout");

            // Confirm retry deadline remains set
            let fourth = state.next_timeout_deadline();
            assert_eq!(fourth, third, "retry deadline should be set");

            // Confirm works if later is far in the future
            context.sleep(Duration::from_secs(10)).await;
            let fifth = state.next_timeout_deadline();
            assert_eq!(fifth, later + retry, "retry deadline should be set");

            // Handle timeout should return true whenever called (can be before registered deadline)
            let (was_retry, _, _) = state.handle_timeout();
            assert!(was_retry, "subsequent timeout should be treated as retry");

            // Confirm retry deadline is set
            let sixth = state.next_timeout_deadline();
            let later = context.current();
            assert_eq!(sixth, later + retry, "retry deadline should be set");
        });
    }

    #[test]
    fn round_prunes_with_min_active() {
        let runtime = deterministic::Runner::default();
        runtime.start(|mut context| async move {
            let namespace = b"ns".to_vec();
            let Fixture {
                schemes, verifier, ..
            } = ed25519::fixture(&mut context, &namespace, 4);
            let cfg = Config {
                scheme: schemes[0].clone(),
                elector: <RoundRobin>::default(),
                epoch: Epoch::new(7),
                activity_timeout: ViewDelta::new(10),
                leader_timeout: Duration::from_secs(1),
                notarization_timeout: Duration::from_secs(2),
                nullify_retry: Duration::from_secs(3),
            };
            let mut state = State::new(context, cfg);
            state.set_genesis(test_genesis());

            // Add initial rounds
            for view in 0..5 {
                state.create_round(View::new(view));
            }

            // Create finalization for view 20
            let proposal_a = Proposal::new(
                Rnd::new(Epoch::new(1), View::new(20)),
                GENESIS_VIEW,
                Sha256Digest::from([1u8; 32]),
            );
            let finalization_votes: Vec<_> = schemes
                .iter()
                .map(|scheme| Finalize::sign(scheme, proposal_a.clone()).unwrap())
                .collect();
            let finalization =
                Finalization::from_finalizes(&verifier, finalization_votes.iter(), &Sequential)
                    .expect("finalization");
            state.add_finalization(finalization);

            // Update last finalize to be in the future
            let removed = state.prune();
            assert_eq!(
                removed,
                vec![
                    View::new(0),
                    View::new(1),
                    View::new(2),
                    View::new(3),
                    View::new(4)
                ]
            );
            assert_eq!(state.views.len(), 2); // 20 and 21
        });
    }

    #[test]
    fn parent_payload_returns_parent_digest() {
        let runtime = deterministic::Runner::default();
        runtime.start(|mut context| async move {
            let namespace = b"ns".to_vec();
            let Fixture {
                schemes, verifier, ..
            } = ed25519::fixture(&mut context, &namespace, 4);
            let local_scheme = schemes[2].clone(); // leader of view 1
            let cfg = Config {
                scheme: local_scheme,
                elector: <RoundRobin>::default(),
                epoch: Epoch::new(4),
                activity_timeout: ViewDelta::new(2),
                leader_timeout: Duration::from_secs(1),
                notarization_timeout: Duration::from_secs(2),
                nullify_retry: Duration::from_secs(3),
            };
            let mut state = State::new(context, cfg);
            state.set_genesis(test_genesis());

            // Create proposal
            let parent_view = View::new(1);
            let parent_payload = Sha256Digest::from([1u8; 32]);
            let parent_proposal = Proposal::new(
                Rnd::new(Epoch::new(1), parent_view),
                GENESIS_VIEW,
                parent_payload,
            );

            // Attempt to get parent payload without certificate
            let proposal = Proposal::new(
                Rnd::new(Epoch::new(1), View::new(2)),
                parent_view,
                Sha256Digest::from([9u8; 32]),
            );
            assert!(state.parent_payload(&proposal).is_none());

            // Add notarization certificate
            let notarization_votes: Vec<_> = schemes
                .iter()
                .map(|scheme| Notarize::sign(scheme, parent_proposal.clone()).unwrap())
                .collect();
            let notarization =
                Notarization::from_notarizes(&verifier, notarization_votes.iter(), &Sequential)
                    .unwrap();
            state.add_notarization(notarization);

            // The parent is still not certified
            assert!(state.parent_payload(&proposal).is_none());

            // Set certify handle then certify the parent
            let mut pool = AbortablePool::<()>::default();
            let handle = pool.push(futures::future::pending());
            state.set_certify_handle(parent_view, handle);
            state.certified(parent_view, true);
            let digest = state.parent_payload(&proposal).expect("parent payload");
            assert_eq!(digest, parent_payload);
        });
    }

    #[test]
    fn parent_certificate_prefers_finalization() {
        let runtime = deterministic::Runner::default();
        runtime.start(|mut context| async move {
            let namespace = b"ns".to_vec();
            let Fixture {
                schemes, verifier, ..
            } = ed25519::fixture(&mut context, &namespace, 4);
            let local_scheme = schemes[1].clone(); // leader of view 2
            let cfg = Config {
                scheme: local_scheme,
                elector: <RoundRobin>::default(),
                epoch: Epoch::new(7),
                activity_timeout: ViewDelta::new(3),
                leader_timeout: Duration::from_secs(1),
                notarization_timeout: Duration::from_secs(2),
                nullify_retry: Duration::from_secs(3),
            };
            let mut state = State::new(context, cfg);
            state.set_genesis(test_genesis());

            // Add notarization for parent view
            let parent_round = Rnd::new(state.epoch(), View::new(1));
            let parent_proposal =
                Proposal::new(parent_round, GENESIS_VIEW, Sha256Digest::from([11u8; 32]));
            let notarize_votes: Vec<_> = schemes
                .iter()
                .map(|scheme| Notarize::sign(scheme, parent_proposal.clone()).unwrap())
                .collect();
            let notarization =
                Notarization::from_notarizes(&verifier, notarize_votes.iter(), &Sequential)
                    .expect("notarization");
            state.add_notarization(notarization.clone());

            // Insert proposal at view 2 with parent at view 1
            let proposal = Proposal::new(
                Rnd::new(state.epoch(), View::new(2)),
                View::new(1),
                Sha256Digest::from([22u8; 32]),
            );
            state.proposed(proposal);

            // parent_certificate returns the notarization
            let cert = state.parent_certificate(View::new(2)).unwrap();
            assert!(matches!(cert, Certificate::Notarization(n) if n == notarization));

            // Add finalization for the same parent view
            let finalize_votes: Vec<_> = schemes
                .iter()
                .map(|scheme| Finalize::sign(scheme, parent_proposal.clone()).unwrap())
                .collect();
            let finalization =
                Finalization::from_finalizes(&verifier, finalize_votes.iter(), &Sequential)
                    .expect("finalization");
            state.add_finalization(finalization.clone());

            // parent_certificate now returns the finalization (preferred)
            let cert = state.parent_certificate(View::new(2)).unwrap();
            assert!(matches!(cert, Certificate::Finalization(f) if f == finalization));
        });
    }

    #[test]
    fn parent_payload_errors_without_nullification() {
        let runtime = deterministic::Runner::default();
        runtime.start(|mut context| async move {
            let namespace = b"ns".to_vec();
            let Fixture {
                schemes, verifier, ..
            } = ed25519::fixture(&mut context, &namespace, 4);
            let cfg = Config {
                scheme: verifier.clone(),
                elector: <RoundRobin>::default(),
                epoch: Epoch::new(1),
                leader_timeout: Duration::from_secs(1),
                notarization_timeout: Duration::from_secs(2),
                nullify_retry: Duration::from_secs(3),
                activity_timeout: ViewDelta::new(5),
            };
            let mut state = State::new(context, cfg);
            state.set_genesis(test_genesis());

            // Create parent proposal and certificate
            let parent_view = View::new(1);
            let parent_proposal = Proposal::new(
                Rnd::new(Epoch::new(1), parent_view),
                GENESIS_VIEW,
                Sha256Digest::from([2u8; 32]),
            );
            let notarization_votes: Vec<_> = schemes
                .iter()
                .map(|scheme| Notarize::sign(scheme, parent_proposal.clone()).unwrap())
                .collect();
            let notarization =
                Notarization::from_notarizes(&verifier, notarization_votes.iter(), &Sequential)
                    .unwrap();
            state.add_notarization(notarization);
            state.create_round(View::new(2));

            // Attempt to get parent payload
            let proposal = Proposal::new(
                Rnd::new(Epoch::new(1), View::new(3)),
                parent_view,
                Sha256Digest::from([3u8; 32]),
            );
            assert!(state.parent_payload(&proposal).is_none());
        });
    }

    #[test]
    fn parent_payload_returns_genesis_payload() {
        let runtime = deterministic::Runner::default();
        runtime.start(|mut context| async move {
            let namespace = b"ns".to_vec();
            let Fixture {
                schemes, verifier, ..
            } = ed25519::fixture(&mut context, &namespace, 4);
            let cfg = Config {
                scheme: verifier.clone(),
                elector: <RoundRobin>::default(),
                epoch: Epoch::new(1),
                leader_timeout: Duration::from_secs(1),
                notarization_timeout: Duration::from_secs(2),
                nullify_retry: Duration::from_secs(3),
                activity_timeout: ViewDelta::new(5),
            };
            let mut state = State::new(context, cfg);
            state.set_genesis(test_genesis());

            // Add nullification certificate for view 1
            let nullify_votes: Vec<_> = schemes
                .iter()
                .map(|scheme| {
                    Nullify::sign::<Sha256Digest>(scheme, Rnd::new(Epoch::new(1), View::new(1)))
                        .unwrap()
                })
                .collect();
            let nullification =
                Nullification::from_nullifies(&verifier, &nullify_votes, &Sequential).unwrap();
            state.add_nullification(nullification);

            // Get genesis payload
            let proposal = Proposal::new(
                Rnd::new(Epoch::new(1), View::new(2)),
                GENESIS_VIEW,
                Sha256Digest::from([8u8; 32]),
            );
            let genesis = Sha256Digest::from([0u8; 32]);
            let digest = state.parent_payload(&proposal).expect("genesis payload");
            assert_eq!(digest, genesis);
        });
    }

    #[test]
    fn parent_payload_rejects_parent_before_finalized() {
        let runtime = deterministic::Runner::default();
        runtime.start(|mut context| async move {
            let namespace = b"ns".to_vec();
            let Fixture {
                schemes, verifier, ..
            } = ed25519::fixture(&mut context, &namespace, 4);
            let cfg = Config {
                scheme: verifier.clone(),
                elector: <RoundRobin>::default(),
                epoch: Epoch::new(1),
                activity_timeout: ViewDelta::new(5),
                leader_timeout: Duration::from_secs(1),
                notarization_timeout: Duration::from_secs(2),
                nullify_retry: Duration::from_secs(3),
            };
            let mut state = State::new(context, cfg);
            state.set_genesis(test_genesis());

            // Add finalization
            let proposal_a = Proposal::new(
                Rnd::new(Epoch::new(1), View::new(3)),
                GENESIS_VIEW,
                Sha256Digest::from([1u8; 32]),
            );
            let finalization_votes: Vec<_> = schemes
                .iter()
                .map(|scheme| Finalize::sign(scheme, proposal_a.clone()).unwrap())
                .collect();
            let finalization =
                Finalization::from_finalizes(&verifier, finalization_votes.iter(), &Sequential)
                    .expect("finalization");
            state.add_finalization(finalization);

            // Attempt to verify before finalized
            let proposal = Proposal::new(
                Rnd::new(Epoch::new(1), View::new(4)),
                View::new(2),
                Sha256Digest::from([6u8; 32]),
            );
            assert!(state.parent_payload(&proposal).is_none());
        });
    }

    #[test]
    fn replay_restores_conflict_state() {
        let runtime = deterministic::Runner::default();
        runtime.start(|mut context| async move {
            let namespace = b"ns".to_vec();
            let Fixture {
                schemes, verifier, ..
            } = ed25519::fixture(&mut context, &namespace, 4);
            let mut scheme_iter = schemes.into_iter();
            let local_scheme = scheme_iter.next().unwrap();
            let other_schemes: Vec<_> = scheme_iter.collect();
            let epoch: Epoch = Epoch::new(3);
            let mut state = State::new(
                context.clone(),
                Config {
                    scheme: local_scheme.clone(),
                    elector: <RoundRobin>::default(),
                    epoch: Epoch::new(1),
                    activity_timeout: ViewDelta::new(5),
                    leader_timeout: Duration::from_secs(1),
                    notarization_timeout: Duration::from_secs(2),
                    nullify_retry: Duration::from_secs(3),
                },
            );
            state.set_genesis(test_genesis());
            let view = View::new(4);
            let round = Rnd::new(epoch, view);
            let proposal_a = Proposal::new(round, GENESIS_VIEW, Sha256Digest::from([21u8; 32]));
            let proposal_b = Proposal::new(round, GENESIS_VIEW, Sha256Digest::from([22u8; 32]));
            let local_vote = Notarize::sign(&local_scheme, proposal_a).unwrap();

            // Replay local notarize vote
            state.replay(&Artifact::Notarize(local_vote.clone()));

            // Add conflicting notarization certificate and replay
            let votes_b: Vec<_> = other_schemes
                .iter()
                .take(3)
                .map(|scheme| Notarize::sign(scheme, proposal_b.clone()).unwrap())
                .collect();
            let conflicting = Notarization::from_notarizes(&verifier, votes_b.iter(), &Sequential)
                .expect("certificate");
            state.add_notarization(conflicting.clone());
            state.replay(&Artifact::Notarization(conflicting.clone()));

            // Shouldn't finalize the certificate's proposal (proposal_b)
            assert!(state.construct_finalize(view).is_none());

            // Restart state and replay
            let mut restarted = State::new(
                context,
                Config {
                    scheme: local_scheme,
                    elector: <RoundRobin>::default(),
                    epoch: Epoch::new(1),
                    activity_timeout: ViewDelta::new(5),
                    leader_timeout: Duration::from_secs(1),
                    notarization_timeout: Duration::from_secs(2),
                    nullify_retry: Duration::from_secs(3),
                },
            );
            restarted.set_genesis(test_genesis());
            restarted.replay(&Artifact::Notarize(local_vote));
            restarted.add_notarization(conflicting.clone());
            restarted.replay(&Artifact::Notarization(conflicting));

            // Shouldn't finalize the certificate's proposal (proposal_b)
            assert!(restarted.construct_finalize(view).is_none());
        });
    }

    #[test]
    fn certification_lifecycle() {
        let runtime = deterministic::Runner::default();
        runtime.start(|mut context| async move {
            let namespace = b"ns".to_vec();
            let Fixture {
                schemes, verifier, ..
            } = ed25519::fixture(&mut context, &namespace, 4);
            let cfg = Config {
                scheme: verifier.clone(),
                elector: <RoundRobin>::default(),
                epoch: Epoch::new(1),
                activity_timeout: ViewDelta::new(10),
                leader_timeout: Duration::from_secs(1),
                notarization_timeout: Duration::from_secs(2),
                nullify_retry: Duration::from_secs(3),
            };
            let mut state = State::new(context, cfg);
            state.set_genesis(test_genesis());

            // Helper to create notarization for a view
            let make_notarization = |view: View| {
                let proposal = Proposal::new(
                    Rnd::new(Epoch::new(1), view),
                    GENESIS_VIEW,
                    Sha256Digest::from([view.get() as u8; 32]),
                );
                let votes: Vec<_> = schemes
                    .iter()
                    .map(|s| Notarize::sign(s, proposal.clone()).unwrap())
                    .collect();
                Notarization::from_notarizes(&verifier, votes.iter(), &Sequential).unwrap()
            };

            // Helper to create finalization for a view
            let make_finalization = |view: View| {
                let proposal = Proposal::new(
                    Rnd::new(Epoch::new(1), view),
                    GENESIS_VIEW,
                    Sha256Digest::from([view.get() as u8; 32]),
                );
                let votes: Vec<_> = schemes
                    .iter()
                    .map(|s| Finalize::sign(s, proposal.clone()).unwrap())
                    .collect();
                Finalization::from_finalizes(&verifier, votes.iter(), &Sequential).unwrap()
            };

            let mut pool = AbortablePool::<()>::default();

            // Add notarizations for views 3-8
            for i in 3..=8u64 {
                state.add_notarization(make_notarization(View::new(i)));
            }

            // All 6 views should be candidates
            let candidates = state.certify_candidates();
            assert_eq!(candidates.len(), 6);

            // Set certify handles for views 3, 4, 5, 7 (NOT 6 or 8)
            for i in [3u64, 4, 5, 7] {
                let handle = pool.push(futures::future::pending());
                state.set_certify_handle(View::new(i), handle);
            }

            // Candidates empty (consumed by certify_candidates, handles block re-fetching)
            assert!(state.certify_candidates().is_empty());

            // Complete certification for view 7 (success)
            let notarization = state.certified(View::new(7), true);
            assert!(notarization.is_some());

            // View 7 should not be aborted (it was certified successfully)
            assert!(!state.is_certify_aborted(View::new(7)));

            // Add finalization for view 5 - aborts handles for views 3, 4, 5
            state.add_finalization(make_finalization(View::new(5)));

            // Verify views 3, 4, 5 had their certification aborted
            assert!(state.is_certify_aborted(View::new(3)));
            assert!(state.is_certify_aborted(View::new(4)));
            assert!(state.is_certify_aborted(View::new(5)));

            // View 7 still not aborted (was certified, and 7 > 5)
            assert!(!state.is_certify_aborted(View::new(7)));

            // Views 6, 8 never had handles set, so they're not aborted (still Ready)
            assert!(!state.is_certify_aborted(View::new(6)));
            assert!(!state.is_certify_aborted(View::new(8)));

            // Candidates empty: 3-5 finalized, 6/8 consumed, 7 certified
            assert!(state.certify_candidates().is_empty());

            // Add view 9, should be returned as candidate
            state.add_notarization(make_notarization(View::new(9)));
            let candidates = state.certify_candidates();
            assert_eq!(candidates.len(), 1);
            assert_eq!(candidates[0].round.view(), View::new(9));

            // Set handle for view 9, add view 10
            let handle9 = pool.push(futures::future::pending());
            state.set_certify_handle(View::new(9), handle9);
            state.add_notarization(make_notarization(View::new(10)));

            // View 10 returned (view 9 has handle)
            let candidates = state.certify_candidates();
            assert_eq!(candidates.len(), 1);
            assert_eq!(candidates[0].round.view(), View::new(10));

            // Finalize view 9 - aborts view 9's handle
            state.add_finalization(make_finalization(View::new(9)));
            assert!(state.is_certify_aborted(View::new(9)));

            // Add view 11, should be returned
            state.add_notarization(make_notarization(View::new(11)));
            let candidates = state.certify_candidates();
            assert_eq!(candidates.len(), 1);
            assert_eq!(candidates[0].round.view(), View::new(11));
        });
    }

    #[test]
    fn only_notarize_before_nullify() {
        let runtime = deterministic::Runner::default();
        runtime.start(|mut context| async move {
            let namespace = b"ns".to_vec();
            let Fixture { schemes, .. } = ed25519::fixture(&mut context, &namespace, 4);
            let cfg = Config {
                scheme: schemes[0].clone(),
                elector: <RoundRobin>::default(),
                epoch: Epoch::new(1),
                activity_timeout: ViewDelta::new(5),
                leader_timeout: Duration::from_secs(1),
                notarization_timeout: Duration::from_secs(2),
                nullify_retry: Duration::from_secs(3),
            };
            let mut state = State::new(context, cfg);
            state.set_genesis(test_genesis());
            let view = state.current_view();

            // Set proposal
            let proposal = Proposal::new(
                Rnd::new(Epoch::new(1), view),
                GENESIS_VIEW,
                Sha256Digest::from([1u8; 32]),
            );
            state.set_proposal(view, proposal);

            // We should not want to verify (already timeout)
            assert!(state.try_verify().is_some());
            assert!(state.verified(view));

            // Handle timeout
            assert!(!state.handle_timeout().0);

            // Attempt to notarize after timeout
            assert!(state.construct_notarize(view).is_none());
        });
    }
}