use super::FinalizedTip;
use crate::dkg::{
    ParticipantsProvider, Registrar, ReshareBlock, SecretStore,
    fence::Fence,
    network::{Directory, Manager},
    reshare::{
        Actor,
        metrics::Phase,
        store::{Dealer, Player, Store},
    },
    types::{EpochInfo, EpochOutcome, Participants, Payload, SchemeInfo},
};
use commonware_consensus::{
    Block,
    marshal::{
        Identifier,
        core::{Processed, Variant as MarshalVariant},
    },
    simplex::scheme::Scheme as SimplexScheme,
    types::{Epoch, EpochPhase, Epocher, FixedEpocher, Height},
};
use commonware_cryptography::{
    BatchVerifier, PublicKey, Signer,
    bls12381::{
        dkg::feldman_desmedt::{Info, Output},
        primitives::{group::Share, variant::Variant as BlsVariant},
    },
    transcript::Summary,
};
use commonware_p2p::Blocker;
use commonware_parallel::Strategy;
use commonware_runtime::{Clock, Metrics, Spawner, Storage};
use commonware_utils::{Acknowledgement, N3f1, ordered::Quorum as _};
use rand_core::CryptoRng;
use tracing::warn;

pub(super) struct PreparedEpoch<V, C>
where
    V: BlsVariant,
    C: Signer,
{
    pub(super) epoch: Epoch,
    pub(super) phase: EpochPhase,
    pub(super) info: Info<V, C::PublicKey>,
    pub(super) dealer: Option<Dealer<V, C>>,
    pub(super) player: Option<Player<V, C>>,
}

pub(super) struct EpochPreparation<V, P>
where
    V: BlsVariant,
    P: PublicKey,
{
    pub(super) epoch: Epoch,
    pub(super) phase: EpochPhase,
    pub(super) participants: Participants<P>,
    pub(super) previous: Option<Output<V, P>>,
    pub(super) share: Option<Share>,
    pub(super) seed: Summary,
}

pub(super) enum Setup<V, C>
where
    V: BlsVariant,
    C: Signer,
{
    Follow,
    Participate(Box<PreparedEpoch<V, C>>),
}

/// State-sync epoch metadata paired with the block height of its certified floor.
pub(super) struct StateSyncStart<V, P, D>
where
    V: BlsVariant,
    P: PublicKey,
    D: Directory<P>,
{
    pub(super) info: EpochInfo<V, P, D>,
    pub(super) floor: Height,
}

/// Returns whether a state-sync floor lies after its epoch's midpoint, so the
/// blocks it skipped may have carried dealer logs.
fn state_sync_skips_inclusion_prefix(
    epocher: &FixedEpocher,
    state_sync_floor: Option<Height>,
) -> bool {
    let Some(floor) = state_sync_floor else {
        return false;
    };
    epocher
        .containing(floor)
        .expect("epocher must know state sync floor")
        .phase()
        == EpochPhase::Late
}

/// Returns `public_key`'s share of `info.output`.
///
/// The share stored for `info.epoch` is returned when present. A failed
/// ceremony carries the previous output forward, so when none is stored for a
/// failed epoch, the search walks back to the newest share stored under an
/// earlier epoch and returns it only if the output's public polynomial confirms
/// it as `public_key`'s share. A node that ran through every boundary stores
/// the carried share under each epoch, while a node that state synced past
/// several failed ceremonies holds it under the last epoch it applied.
async fn recovered_share<E, SS, V, P, D>(
    store: &mut Store<E, SS, V, P, D>,
    public_key: &P,
    info: &EpochInfo<V, P, D>,
) -> Option<Share>
where
    E: Clock + Storage + Metrics,
    SS: SecretStore,
    V: BlsVariant,
    P: PublicKey,
    D: Directory<P>,
{
    let share = store.share(info.epoch).await;
    if share.is_some() || info.outcome != EpochOutcome::Failure {
        return share;
    }

    let index = info.output.players().index(public_key)?;
    let expected = info.output.public().partial_public(index).ok()?;
    let mut epoch = info.epoch;
    while let Some(previous) = epoch.previous() {
        if let Some(share) = store.share(previous).await {
            return (share.public::<V>() == expected).then_some(share);
        }
        epoch = previous;
    }
    None
}

/// Registers `scheme_info` for `epoch` and then marks `fence`, so a gate
/// waiting on `epoch` opens only once the consensus provider serves the scheme.
async fn register_and_mark<R: Registrar>(
    registrar: &R,
    fence: &Fence,
    epoch: Epoch,
    scheme_info: SchemeInfo<R::Variant, R::PublicKey>,
) {
    registrar.register(epoch, scheme_info).await;
    fence.mark(epoch);
}

/// Returns the later of marshal's processed position and the parent of the
/// state-sync floor (`None` when neither exists).
///
/// `processed` pairs the position with the ready processed block or its child
/// when the processed block is absent.
pub(super) fn startup<B: Block>(
    processed: Option<(Processed, B)>,
    floor: Option<B>,
) -> Option<FinalizedTip<B::Digest>> {
    let processed = processed.map(|(processed, block)| match processed {
        Processed::Block(height) => FinalizedTip {
            height,
            digest: block.digest(),
        },
        Processed::Absent(height) => FinalizedTip {
            height,
            digest: block.parent(),
        },
    });
    let floor = floor.and_then(|block| {
        Some(FinalizedTip {
            height: block.height().previous()?,
            digest: block.parent(),
        })
    });
    processed
        .into_iter()
        .chain(floor)
        .max_by_key(|tip| tip.height)
}

/// Retains a dealer that can still produce a selectable log.
///
/// During the early phase, missing acknowledgements can still arrive. Afterward,
/// a recovered dealer is useful only when its durable acknowledgements already
/// form a quorum.
fn dealer_for_phase<V, C>(
    phase: EpochPhase,
    players: usize,
    dealer: Option<Dealer<V, C>>,
) -> Option<Dealer<V, C>>
where
    V: BlsVariant,
    C: Signer,
{
    dealer.filter(|dealer| {
        phase == EpochPhase::Early || dealer.has_acknowledgement_quorum::<N3f1>(players)
    })
}

impl<E, B, V, C, M, X, P, SS, T, S, MV, R, A> Actor<E, B, V, C, M, X, P, SS, T, S, MV, R, A>
where
    E: Spawner + CryptoRng + Metrics + Clock + Storage,
    B: ReshareBlock<Variant = V, PublicKey = C::PublicKey>,
    V: BlsVariant,
    C: Signer,
    M: Manager<PublicKey = C::PublicKey, Directory = B::Directory>,
    X: Blocker<PublicKey = C::PublicKey>,
    P: ParticipantsProvider<PublicKey = C::PublicKey, Directory = B::Directory>,
    SS: SecretStore,
    T: Strategy,
    C::PublicKey: BatchVerifier,
    S: SimplexScheme<MV::Commitment, PublicKey = C::PublicKey>,
    MV: MarshalVariant<ApplicationBlock = B>,
    R: Registrar<Variant = V, PublicKey = C::PublicKey>,
    A: Acknowledgement,
{
    pub(super) async fn setup(
        &mut self,
        store: &mut Store<E, SS, V, C::PublicKey, B::Directory>,
        state_sync: Option<StateSyncStart<V, C::PublicKey, B::Directory>>,
    ) -> Setup<V, C> {
        self.metrics.set_phase(Phase::Setup);

        let height = self.tip.map_or_else(Height::zero, |tip| tip.height.next());
        let bounds = self
            .epocher
            .containing(height)
            .expect("epocher must know of block height");
        let epoch = bounds.epoch();
        self.metrics.set_epoch(epoch.get());

        // Prefer the epoch committed earlier in this run, then state-sync
        // material, then the finalized boundary block.
        let current = store
            .current()
            .filter(|current| current.epoch == epoch)
            .cloned();
        let already_committed = current.is_some();
        let state_sync_floor = state_sync.as_ref().map(|start| start.floor);
        let follow = state_sync_skips_inclusion_prefix(&self.epocher, state_sync_floor);
        let info = match current.or_else(|| state_sync.map(|start| start.info)) {
            Some(info) => info,
            None => {
                let Some(info) = self.boundary_epoch_info(epoch).await else {
                    return Setup::Follow;
                };
                info
            }
        };
        if info.epoch != epoch {
            panic!(
                "boundary epoch info describes epoch {}, expected {epoch}",
                info.epoch
            );
        }

        // Each source is read from storage or produced by consensus, so check
        // participant limits before building protocol state.
        let participants = info.participants();
        let round = epoch.get();
        participants
            .validate(self.max_participants, Some(&info.output), round)
            .expect("boundary epoch participants must be valid");

        // Persist the epoch's share and seed before registering its scheme, and
        // prune only once the epoch is durable.
        let share = self.recovered_share(store, &info).await;
        let seed = store
            .seed_or_random(epoch, self.context.as_present_mut())
            .await;
        if !already_committed {
            store.commit_epoch(info.clone(), seed, share.clone()).await;
            self.register_epoch(&info, share.clone()).await;
        }
        store.prune(epoch.previous().unwrap_or(epoch)).await;

        // A floor after the midpoint skipped blocks that may carry dealer logs,
        // so follow the committed epoch without participating. A follower that
        // holds a share still votes in the epoch's consensus but cannot verify
        // its final block, so it counts as a faulty validator at that block.
        if follow {
            if share.is_some() {
                warn!(
                    ?epoch,
                    floor = ?state_sync_floor,
                    "validator synced past the epoch midpoint, cannot vote to finalize the epoch's last block"
                );
            }
            return Setup::Follow;
        }

        Setup::Participate(Box::new(self.prepare_epoch(
            store,
            EpochPreparation {
                epoch,
                phase: bounds.phase(),
                participants,
                previous: Some(info.output.clone()),
                share,
                seed,
            },
        )))
    }

    /// Returns this node's share of `info.output`, found by [`recovered_share`].
    pub(super) async fn recovered_share(
        &mut self,
        store: &mut Store<E, SS, V, C::PublicKey, B::Directory>,
        info: &EpochInfo<V, C::PublicKey, B::Directory>,
    ) -> Option<Share> {
        recovered_share(store, &self.signer.public_key(), info).await
    }

    /// Returns the [`EpochInfo`] carried by the finalized boundary block that
    /// introduced `epoch` (`None` if marshal does not have the block).
    ///
    /// Panics if the block carries no [`EpochInfo`].
    async fn boundary_epoch_info(
        &mut self,
        epoch: Epoch,
    ) -> Option<EpochInfo<V, C::PublicKey, B::Directory>> {
        let height = epoch
            .previous()
            .and_then(|e| self.epocher.last(e))
            .unwrap_or(Height::zero());
        let block = self
            .marshal
            .get_block(Identifier::Height(height))
            .await
            .map(MV::into_shared)?;
        let Some(Payload::EpochInfo(info)) = block.inner().payload() else {
            panic!("boundary block {height} missing epoch info");
        };
        Some(info)
    }

    /// Registers the scheme for `info` in this node's role, then marks `info.epoch` on the
    /// fence.
    pub(super) async fn register_epoch(
        &mut self,
        info: &EpochInfo<V, C::PublicKey, B::Directory>,
        share: Option<Share>,
    ) {
        let scheme_info = SchemeInfo::new(
            &self.signer.public_key(),
            info.output.players().clone(),
            info.output.public().clone(),
            share,
        );
        register_and_mark(&self.registrar, &self.fence, info.epoch, scheme_info).await;
    }

    pub(super) fn prepare_epoch(
        &mut self,
        store: &mut Store<E, SS, V, C::PublicKey, B::Directory>,
        preparation: EpochPreparation<V, C::PublicKey>,
    ) -> PreparedEpoch<V, C> {
        let EpochPreparation {
            epoch,
            phase,
            participants,
            previous,
            share,
            seed,
        } = preparation;

        let round = Info::new::<N3f1>(
            self.namespace,
            epoch.get(),
            previous.clone(),
            self.sharing_mode,
            self.reveal,
            participants.dealers.clone(),
            participants.players.clone(),
        )
        .expect("epoch participants must produce valid round info");

        // Reshare dealers need the prior private share, while epoch-zero DKG has
        // no prior output. Player construction depends only on membership in the
        // new player set.
        let public_key = self.signer.public_key();
        let has_prior_share = previous.is_none() || share.is_some();
        let dealer = if participants.dealers.position(&public_key).is_some() && has_prior_share {
            store.create_dealer::<C, N3f1>(epoch, self.signer.clone(), round.clone(), share, seed)
        } else {
            None
        };
        let dealer = dealer_for_phase(phase, participants.players.len(), dealer);
        let player = participants.players.position(&public_key).and_then(|_| {
            store.create_player::<C, N3f1>(epoch, self.signer.clone(), round.clone())
        });

        PreparedEpoch {
            epoch,
            phase,
            info: round,
            dealer,
            player,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{
        dealer_for_phase, recovered_share, register_and_mark, startup,
        state_sync_skips_inclusion_prefix,
    };
    use crate::dkg::{
        Registrar, SecretStore as _,
        fence::Fence,
        reshare::store::{AckOutcome, Store},
        tests::mocks::{self, MemorySecretStore, TestBlsVariant, child, genesis_block},
        types::{EpochInfo, EpochOutcome, SchemeInfo},
    };
    use commonware_consensus::{
        marshal::core::Processed,
        types::{Epoch, EpochPhase, Epocher as _, FixedEpocher, Height},
    };
    use commonware_cryptography::{
        Digested, Digestible as _, Signer,
        bls12381::{
            dkg::feldman_desmedt::{Info, Player, Reveal, deal},
            primitives::sharing::Mode,
        },
        ed25519::{PrivateKey, PublicKey},
    };
    use commonware_runtime::{Runner, Supervisor as _, deterministic};
    use commonware_utils::{
        N3f1, NZU64, TestRng, channel::oneshot, ordered::Set, sequence::Unit, sync::Mutex, test_rng,
    };
    use futures::poll;
    use std::pin::pin;

    const TEST_NAMESPACE: &[u8] = b"_COMMONWARE_GLUE_DKG_RESHARE_SETUP_TEST";

    /// State sync never resumes below its certified floor, and marshal's
    /// processed position takes its digest from the stored block that backs it.
    #[test]
    fn startup_tip_covers_processed_and_floor() {
        let mut blocks = vec![genesis_block(PrivateKey::from_seed(0).public_key())];
        while blocks.len() < 5 {
            blocks.push(child(blocks.last().unwrap()));
        }
        let tip = |processed: Option<Processed>, floor: Option<usize>| {
            let processed = processed.map(|processed| {
                let anchor = processed.anchor().get() as usize;
                (
                    processed,
                    Digested::precomputed(blocks[anchor].clone(), blocks[anchor].digest()),
                )
            });
            startup(
                processed,
                floor.map(|floor| {
                    Digested::precomputed(blocks[floor].clone(), blocks[floor].digest())
                }),
            )
            .map(|tip| (tip.height, tip.digest))
        };
        let covered = |index: usize| Some((Height::new(index as u64), blocks[index].digest()));
        let stored = |height| Some(Processed::Block(Height::new(height)));
        let absent = |height| Some(Processed::Absent(Height::new(height)));

        // Without state sync, the actor resumes after marshal's processed position.
        assert_eq!(tip(None, None), None);
        assert_eq!(tip(stored(2), None), covered(2));
        assert_eq!(tip(absent(2), None), covered(2));

        // A floor covers every block below it until marshal passes it.
        assert_eq!(tip(None, Some(0)), None);
        assert_eq!(tip(None, Some(3)), covered(2));
        assert_eq!(tip(stored(1), Some(3)), covered(2));
        assert_eq!(tip(absent(2), Some(3)), covered(2));
        assert_eq!(tip(stored(4), Some(3)), covered(4));
    }

    #[test]
    fn late_state_sync_floor_skips_inclusion_prefix() {
        let epocher = FixedEpocher::new(NZU64!(64));
        let epoch = Epoch::new(3);
        let midpoint = epocher.midpoint(epoch).expect("test epoch");
        let late_floor = midpoint.next();
        let late_phase = epocher.containing(late_floor).expect("test floor").phase();

        assert_eq!(late_phase, EpochPhase::Late);
        assert!(state_sync_skips_inclusion_prefix(
            &epocher,
            Some(late_floor)
        ));
        assert!(!state_sync_skips_inclusion_prefix(&epocher, Some(midpoint)));
        assert!(!state_sync_skips_inclusion_prefix(&epocher, None));
    }

    /// After a failed ceremony, the newest share stored under an earlier epoch
    /// is recovered only if it is a share of the carried output. A successful
    /// epoch with no stored share recovers nothing, and neither does a node
    /// that is not a player of the carried output.
    #[test]
    fn recovered_share_matches_carried_output() {
        let executor = deterministic::Runner::default();
        executor.start(|context| async move {
            let signers = (0..4).map(PrivateKey::from_seed).collect::<Vec<_>>();
            let players = Set::from_iter_dedup(signers.iter().map(Signer::public_key));
            let (carried, shares) =
                deal::<TestBlsVariant, _, N3f1>(test_rng(), Mode::NonZeroCounter, players.clone())
                    .expect("trusted dealing must succeed");
            let (other, other_shares) = deal::<TestBlsVariant, _, N3f1>(
                TestRng::new(1),
                Mode::NonZeroCounter,
                players.clone(),
            )
            .expect("other trusted dealing must succeed");
            let public_key = signers[0].public_key();
            let share = shares
                .get_value(&public_key)
                .expect("signer must hold a share of the carried output")
                .clone();
            let info = |epoch: u64, outcome: EpochOutcome| EpochInfo {
                outcome,
                epoch: Epoch::new(epoch),
                output: carried.clone(),
                players: players.clone(),
                next_players: players.clone(),
                directory: Unit,
            };
            let mut secrets = MemorySecretStore::default();
            secrets.put_share(Epoch::new(2), share.clone()).await;
            let mut store: Store<_, _, TestBlsVariant, PublicKey> =
                mocks::store(context.child("store"), "recovered-share", secrets.clone()).await;

            // Three failed ceremonies after the stored epoch still carry the
            // stored output, so the share under epoch two is this node's share.
            assert_eq!(
                recovered_share(&mut store, &public_key, &info(5, EpochOutcome::Failure)).await,
                Some(share.clone())
            );

            // A successful ceremony replaces the output, so an unstored epoch
            // has no share to recover even though older shares exist.
            assert_eq!(
                recovered_share(&mut store, &public_key, &info(5, EpochOutcome::Success)).await,
                None
            );

            // A node that is not a player of the carried output recovers
            // nothing, even though a share of that output is stored.
            let outsider = PrivateKey::from_seed(9).public_key();
            assert_eq!(
                recovered_share(&mut store, &outsider, &info(5, EpochOutcome::Failure)).await,
                None
            );

            // The newest stored share decides, so a share of another output
            // stored under epoch three leaves nothing to recover.
            let foreign = other_shares
                .get_value(&public_key)
                .expect("signer must hold a share of the other output")
                .clone();
            assert_ne!(other.public(), carried.public());
            secrets.put_share(Epoch::new(3), foreign).await;
            assert_eq!(
                recovered_share(&mut store, &public_key, &info(5, EpochOutcome::Failure)).await,
                None
            );
        });
    }

    /// The fence marks an epoch only once the registrar resolves, so a gate
    /// waiting on the epoch never opens before the consensus provider serves
    /// its scheme.
    #[test]
    fn fence_marks_after_registration_resolves() {
        struct HeldRegistrar(Mutex<Option<oneshot::Receiver<()>>>);

        impl Registrar for HeldRegistrar {
            type Variant = TestBlsVariant;
            type PublicKey = PublicKey;

            async fn register(&self, _: Epoch, _: SchemeInfo<Self::Variant, Self::PublicKey>) {
                let held = self.0.lock().take().expect("registrar must be called once");
                held.await.expect("test must release the registration");
            }
        }

        let executor = deterministic::Runner::default();
        executor.start(|_| async move {
            let signers = (0..4).map(PrivateKey::from_seed).collect::<Vec<_>>();
            let players = Set::from_iter_dedup(signers.iter().map(Signer::public_key));
            let (output, _) =
                deal::<TestBlsVariant, _, N3f1>(test_rng(), Mode::NonZeroCounter, players.clone())
                    .expect("trusted dealing must succeed");
            let scheme_info = SchemeInfo::Verifier {
                participants: players,
                sharing: output.public().clone(),
            };
            let (release, held) = oneshot::channel();
            let registrar = HeldRegistrar(Mutex::new(Some(held)));
            let (fence, mut gate) = Fence::new();
            let epoch = Epoch::new(1);

            // Registration is suspended inside the registrar, so the gate must
            // still be closed.
            let mut registration = pin!(register_and_mark(&registrar, &fence, epoch, scheme_info));
            assert!(poll!(registration.as_mut()).is_pending());
            assert!(poll!(pin!(gate.wait(epoch))).is_pending());

            // Once released, registration resolves and the gate opens.
            release.send(()).expect("registration must be waiting");
            registration.await;
            gate.wait(epoch).await.expect("fence must still be open");
        });
    }

    #[test]
    fn recovered_dealer_after_dealing_requires_ack_quorum() {
        let executor = deterministic::Runner::default();
        executor.start(|context| async move {
            let epoch = Epoch::new(1);
            let signers = (0..4).map(PrivateKey::from_seed).collect::<Vec<_>>();
            let players = Set::from_iter_dedup(signers.iter().map(Signer::public_key));
            let (previous, shares) =
                deal::<TestBlsVariant, _, N3f1>(test_rng(), Mode::NonZeroCounter, players.clone())
                    .expect("trusted previous output");
            let signer = signers[0].clone();
            let share = shares
                .get_value(&signer.public_key())
                .expect("dealer share")
                .clone();
            let info = Info::new::<N3f1>(
                TEST_NAMESPACE,
                epoch.get(),
                Some(previous),
                Mode::NonZeroCounter,
                Reveal::V1,
                players.clone(),
                players.clone(),
            )
            .expect("valid reshare info");
            let secret_store = MemorySecretStore::default();
            let mut store: Store<_, _, TestBlsVariant, PublicKey> = mocks::store(
                context.child("store"),
                "recovered-dealer-after-dealing",
                secret_store,
            )
            .await;
            let seed = store.seed_or_random(epoch, test_rng()).await;
            let (early_dealer, incomplete_dealer, mut recovered_dealer) = {
                let new_dealer = || {
                    store
                        .create_dealer::<PrivateKey, N3f1>(
                            epoch,
                            signer.clone(),
                            info.clone(),
                            Some(share.clone()),
                            seed,
                        )
                        .expect("current dealer")
                };
                (new_dealer(), new_dealer(), new_dealer())
            };
            assert!(
                dealer_for_phase(EpochPhase::Early, players.len(), Some(early_dealer)).is_some()
            );
            assert!(
                dealer_for_phase(EpochPhase::Midpoint, players.len(), Some(incomplete_dealer))
                    .is_none()
            );

            let dealer_key = signer.public_key();
            for player_signer in &signers[1..] {
                let player_key = player_signer.public_key();
                let mut player =
                    Player::new(info.clone(), player_signer.clone()).expect("current player");
                let (_, public, private) = recovered_dealer
                    .shares_to_distribute()
                    .find(|(recipient, _, _)| recipient == &player_key)
                    .expect("player dealing");
                let ack = player
                    .dealer_message::<N3f1>(dealer_key.clone(), public, private)
                    .expect("valid dealing")
                    .expect("new dealing");
                assert!(matches!(
                    recovered_dealer
                        .handle(&mut store, epoch, player_key, ack)
                        .await,
                    Ok(AckOutcome::Recorded)
                ));
            }
            assert!(
                dealer_for_phase(EpochPhase::Midpoint, players.len(), Some(recovered_dealer))
                    .is_some()
            );
        });
    }
}