//! One-shot engine for generating an initial BLS threshold output.
//!
//! The engine runs an independent Ed25519 Simplex chain for one epoch and uses
//! the reshare actor's crate-private DKG mode to perform the ceremony.
//! The resulting [`EpochInfo`] describes only the ceremony participants. An
//! application using it to start continuous resharing must supply the next
//! players and a transport directory covering the resulting participant union.
//!
//! See [`reshare`] for the protocol flow that this engine reuses and for the
//! application contract of a continuously reshared chain.

use crate::dkg::{
    ParticipantsProvider, Registrar, ReshareBlock, SecretStore,
    fence::Fence,
    network::{Directory, Manager},
    reshare::{self, DkgConfig},
    state_sync::Plan as StateSyncPlan,
    types::{EpochInfo, Participants, Payload, SchemeInfo},
};
use commonware_broadcast::buffered;
use commonware_codec::{Encode, EncodeSize, Error as CodecError, Read, ReadExt as _, Write};
use commonware_consensus::{
    Application, Block as ConsensusBlock, CertifiableBlock, Heightable,
    marshal::{
        self, Start, ancestry::Ancestry, core::Actor as MarshalActor,
        resolver::p2p as marshal_resolver, standard::Deferred,
    },
    simplex::{
        self, Floor,
        config::{ForwardPolicy, SkipBudget, SkipPolicy},
        elector::RoundRobin,
        types::Context,
    },
    types::{Epoch, FixedEpocher, Height, Round, View, ViewDelta},
};
use commonware_cryptography::{
    BatchVerifier, Digest as _, Digestible, Hasher, PublicKey, Sha256, Signer as _,
    bls12381::{
        dkg::feldman_desmedt::Reveal,
        primitives::{
            sharing::{Mode as SharingMode, ModeVersion},
            variant::Variant,
        },
    },
    certificate::{ConstantProvider, Verifier as _},
    ed25519,
    sha256::{self, Digest as Sha256Digest},
};
use commonware_p2p::{Blocker, Receiver, Sender};
use commonware_parallel::Strategy;
use commonware_runtime::{
    Buf, BufMut, BufferPooler, Clock, ContextCell, Handle, Metrics, Spawner, Storage,
    buffer::paged::CacheRef, spawn_cell,
};
use commonware_storage::{archive::prunable, translator::TwoCap};
use commonware_utils::{
    NZU16, NZU32, NZU64, NZUsize,
    channel::{fallible::OneshotExt, oneshot},
    ordered::Set,
    sequence::Unit,
};
use rand_core::{CryptoRng, Rng};
use std::{
    marker::PhantomData,
    num::{NonZeroU16, NonZeroU32, NonZeroU64, NonZeroUsize},
    time::Duration,
};

const MAILBOX_SIZE: NonZeroUsize = NZUsize!(100);
const PAGE_SIZE: NonZeroU16 = NZU16!(1024);
const PAGE_CACHE_PAGES: NonZeroUsize = NZUsize!(16);
const IO_BUFFER_SIZE: NonZeroUsize = NZUsize!(2048);
const ARCHIVE_ITEMS_PER_SECTION: NonZeroU64 = NZU64!(10);

type ConsensusScheme = simplex::scheme::ed25519::Scheme;

/// Configuration for [`Engine`].
pub struct Config<M, X, SS, T, D = Unit> {
    /// Ed25519 signer used for the one-shot consensus chain and DKG protocol messages.
    pub signer: ed25519::PrivateKey,

    /// P2P manager used for peer tracking.
    pub manager: M,

    /// Blocker used for invalid peer behavior.
    pub blocker: X,

    /// User-owned store for private DKG material.
    pub secret_store: SS,

    /// Parallel verification strategy.
    pub strategy: T,

    /// Application namespace for DKG transcript separation.
    pub namespace: &'static [u8],

    /// Sharing mode used for the generated threshold output.
    pub sharing_mode: SharingMode,

    /// Revealed-share calculation used for the DKG ceremony.
    pub reveal: Reveal,

    /// Maximum sharing mode version accepted when decoding blocks.
    pub max_supported_mode: ModeVersion,

    /// Runtime-storage partition prefix.
    pub partition_prefix: String,

    /// Participants in the DKG.
    pub participants: Set<ed25519::PublicKey>,

    /// Transport directory for the participants.
    ///
    /// Used to activate the one-shot chain's peer set and embedded verbatim in
    /// the emitted genesis artifact. Every participant must configure the same
    /// directory.
    pub directory: D,

    /// Length of the one-shot consensus epoch.
    pub blocks_per_epoch: NonZeroU64,
}

/// Completion produced when the one-shot DKG chain finalizes its final block.
pub struct Completion<V: Variant, D: Directory<ed25519::PublicKey> = Unit> {
    /// Final DKG artifact, if the ceremony succeeded.
    ///
    /// Its `next_players` set is empty and its directory covers the one-shot
    /// ceremony participants. Before using it as continuous-resharing genesis,
    /// the application must choose a nonempty next-player set and ensure the
    /// directory exactly covers the resulting participant union.
    pub info: Option<EpochInfo<V, ed25519::PublicKey, D>>,
}

/// Block type used by the one-shot DKG chain.
#[derive(Clone, PartialEq, Eq)]
pub struct Block<V: Variant, D: Directory<ed25519::PublicKey> = Unit> {
    context: Context<sha256::Digest, ed25519::PublicKey>,
    parent: sha256::Digest,
    height: Height,
    payload: Option<Payload<V, ed25519::PrivateKey, D>>,
}

impl<V: Variant, D: Directory<ed25519::PublicKey>> Block<V, D> {
    const fn genesis(leader: ed25519::PublicKey) -> Self {
        Self {
            context: Context {
                round: Round::new(Epoch::zero(), View::zero()),
                leader,
                parent: (View::zero(), Sha256Digest::EMPTY),
            },
            parent: Sha256Digest::EMPTY,
            height: Height::zero(),
            payload: None,
        }
    }

    /// Returns the DKG result carried by this block, if present.
    pub const fn epoch_info(&self) -> Option<&EpochInfo<V, ed25519::PublicKey, D>> {
        match &self.payload {
            Some(Payload::EpochInfo(info)) => Some(info),
            _ => None,
        }
    }
}

impl<V: Variant, D: Directory<ed25519::PublicKey>> Write for Block<V, D> {
    fn write(&self, buf: &mut impl BufMut) {
        self.context.write(buf);
        self.parent.write(buf);
        self.height.write(buf);
        self.payload.write(buf);
    }
}

impl<V: Variant, D: Directory<ed25519::PublicKey>> EncodeSize for Block<V, D> {
    fn encode_size(&self) -> usize {
        self.context.encode_size()
            + self.parent.encode_size()
            + self.height.encode_size()
            + self.payload.encode_size()
    }
}

impl<V: Variant, D: Directory<ed25519::PublicKey>> Read for Block<V, D> {
    type Cfg = (NonZeroU32, ModeVersion);

    fn read_cfg(buf: &mut impl Buf, cfg: &Self::Cfg) -> Result<Self, CodecError> {
        Ok(Self {
            context: Context::read(buf)?,
            parent: sha256::Digest::read(buf)?,
            height: Height::read(buf)?,
            payload: Option::<Payload<V, ed25519::PrivateKey, D>>::read_cfg(buf, cfg)?,
        })
    }
}

impl<V: Variant, D: Directory<ed25519::PublicKey>> Digestible for Block<V, D> {
    type Digest = sha256::Digest;

    fn digest(&self) -> sha256::Digest {
        Sha256::hash(&[&self.encode()])
    }
}

impl<V: Variant, D: Directory<ed25519::PublicKey>> Heightable for Block<V, D> {
    fn height(&self) -> Height {
        self.height
    }
}

impl<V: Variant, D: Directory<ed25519::PublicKey>> ConsensusBlock for Block<V, D> {
    fn parent(&self) -> sha256::Digest {
        self.parent
    }
}

impl<V: Variant, D: Directory<ed25519::PublicKey>> CertifiableBlock for Block<V, D> {
    type Context = Context<sha256::Digest, ed25519::PublicKey>;

    fn context(&self) -> Self::Context {
        self.context.clone()
    }
}

impl<V: Variant, D: Directory<ed25519::PublicKey>> ReshareBlock for Block<V, D> {
    type Variant = V;
    type Signer = ed25519::PrivateKey;
    type Directory = D;

    fn payload(&self) -> Option<Payload<Self::Variant, Self::Signer, Self::Directory>> {
        self.payload.clone()
    }
}

/// Self-contained DKG engine.
pub struct Engine<E, V, M, X, SS, T, D = Unit>
where
    V: Variant,
{
    context: ContextCell<E>,
    config: Config<M, X, SS, T, D>,
    _variant: PhantomData<V>,
}

impl<E, V, M, X, SS, T, D> Engine<E, V, M, X, SS, T, D>
where
    V: Variant,
{
    /// Creates a new engine.
    pub const fn new(context: E, config: Config<M, X, SS, T, D>) -> Self {
        assert!(
            config.max_supported_mode.supports(&config.sharing_mode),
            "sharing mode must be supported by max supported mode",
        );
        Self {
            context: ContextCell::new(context),
            config,
            _variant: PhantomData,
        }
    }
}

impl<E, V, M, X, SS, T, D> Engine<E, V, M, X, SS, T, D>
where
    E: CryptoRng + Spawner + Metrics + Clock + Storage + BufferPooler,
    V: Variant,
    M: Manager<PublicKey = ed25519::PublicKey, Directory = D> + Clone,
    X: Blocker<PublicKey = ed25519::PublicKey> + Clone,
    SS: SecretStore,
    T: Strategy + Clone,
    D: Directory<ed25519::PublicKey>,
    ed25519::Batch: BatchVerifier<PublicKey = ed25519::PublicKey> + Send + 'static,
{
    /// Starts consensus, marshal, broadcast, and the private reshare DKG actor.
    #[allow(clippy::type_complexity, clippy::too_many_arguments)]
    pub fn start(
        mut self,
        votes: (
            impl Sender<PublicKey = ed25519::PublicKey>,
            impl Receiver<PublicKey = ed25519::PublicKey>,
        ),
        certificates: (
            impl Sender<PublicKey = ed25519::PublicKey>,
            impl Receiver<PublicKey = ed25519::PublicKey>,
        ),
        resolver: (
            impl Sender<PublicKey = ed25519::PublicKey>,
            impl Receiver<PublicKey = ed25519::PublicKey>,
        ),
        backfill: (
            impl Sender<PublicKey = ed25519::PublicKey>,
            impl Receiver<PublicKey = ed25519::PublicKey>,
        ),
        broadcast: (
            impl Sender<PublicKey = ed25519::PublicKey>,
            impl Receiver<PublicKey = ed25519::PublicKey>,
        ),
        dkg: (
            impl Sender<PublicKey = ed25519::PublicKey>,
            impl Receiver<PublicKey = ed25519::PublicKey>,
        ),
    ) -> (Handle<()>, oneshot::Receiver<Completion<V, D>>) {
        let (completion_tx, completion_rx) = oneshot::channel();
        let handle = spawn_cell!(
            self.context,
            self.run(
                votes,
                certificates,
                resolver,
                backfill,
                broadcast,
                dkg,
                completion_tx
            )
        );
        (handle, completion_rx)
    }

    #[allow(clippy::too_many_arguments)]
    async fn run(
        self,
        votes: (
            impl Sender<PublicKey = ed25519::PublicKey>,
            impl Receiver<PublicKey = ed25519::PublicKey>,
        ),
        certificates: (
            impl Sender<PublicKey = ed25519::PublicKey>,
            impl Receiver<PublicKey = ed25519::PublicKey>,
        ),
        resolver_network: (
            impl Sender<PublicKey = ed25519::PublicKey>,
            impl Receiver<PublicKey = ed25519::PublicKey>,
        ),
        backfill: (
            impl Sender<PublicKey = ed25519::PublicKey>,
            impl Receiver<PublicKey = ed25519::PublicKey>,
        ),
        broadcast: (
            impl Sender<PublicKey = ed25519::PublicKey>,
            impl Receiver<PublicKey = ed25519::PublicKey>,
        ),
        dkg: (
            impl Sender<PublicKey = ed25519::PublicKey>,
            impl Receiver<PublicKey = ed25519::PublicKey>,
        ),
        completion: oneshot::Sender<Completion<V, D>>,
    ) {
        assert!(
            !self.config.participants.is_empty(),
            "DKG requires at least one participant"
        );
        Participants {
            dealers: self.config.participants.clone(),
            players: self.config.participants.clone(),
            next_players: Set::default(),
        }
        .validate_epoch_capacity::<V>(self.config.blocks_per_epoch, None)
        .expect("DKG epoch must have enough dealer-log slots");
        let participants = self
            .config
            .participants
            .len()
            .try_into()
            .expect("too many DKG participants");
        let max_participants = NZU32!(participants);
        let block_codec_config = (max_participants, self.config.max_supported_mode);

        let context = self.context.into_present();
        let page_cache = CacheRef::from_pooler(&context, PAGE_SIZE, PAGE_CACHE_PAGES);
        let public_key = self.config.signer.public_key();
        let consensus_namespace = [self.config.namespace, b"_INITIAL_CONSENSUS"].concat();
        let scheme = ConsensusScheme::signer(
            &consensus_namespace,
            self.config.participants.clone(),
            self.config.signer.clone(),
        )
        .expect("DKG signer must be a participant");
        let provider = ConstantProvider::<_, Epoch>::new(scheme.clone());
        let genesis = Block::<V, D>::genesis(
            self.config
                .participants
                .iter()
                .next()
                .expect("participants must be non-empty")
                .clone(),
        );

        let (buffer, buffer_mailbox) = buffered::Engine::new(
            context.child("buffer"),
            buffered::Config {
                public_key: public_key.clone(),
                mailbox_size: MAILBOX_SIZE,
                deque_size: 16,
                priority: false,
                codec_config: block_codec_config,
                peer_provider: self.config.manager.clone(),
            },
        );
        let buffer_handle = buffer.start(broadcast);

        let (backfill_handler, backfill_resolver) = marshal_resolver::init(
            context.child("backfill"),
            marshal_resolver::Config {
                public_key: public_key.clone(),
                peer_provider: self.config.manager.clone(),
                blocker: self.config.blocker.clone(),
                mailbox_size: MAILBOX_SIZE,
                timeout: Duration::from_secs(2),
                fetch_retry_timeout: Duration::from_millis(100),
                priority_requests: false,
                priority_responses: false,
            },
            backfill,
        );

        let finalizations = prunable::Archive::init(
            context.child("finalizations"),
            archive_config(
                &self.config.partition_prefix,
                "finalizations",
                page_cache.clone(),
                ConsensusScheme::certificate_codec_config_unbounded(),
            ),
        )
        .await
        .expect("failed to initialize DKG finalization archive");
        let blocks = prunable::Archive::init(
            context.child("blocks"),
            archive_config(
                &self.config.partition_prefix,
                "blocks",
                page_cache.clone(),
                block_codec_config,
            ),
        )
        .await
        .expect("failed to initialize DKG block archive");

        let (marshal_actor, marshal_mailbox, _) = MarshalActor::init(
            context.child("marshal"),
            finalizations,
            blocks,
            marshal::Config {
                provider: provider.clone(),
                epocher: FixedEpocher::new(self.config.blocks_per_epoch),
                start: Start::Genesis(genesis.clone()),
                partition_prefix: format!("{}-marshal", self.config.partition_prefix),
                mailbox_size: MAILBOX_SIZE,
                view_retention: ViewDelta::new(10),
                prunable_items_per_section: ARCHIVE_ITEMS_PER_SECTION,
                page_cache: page_cache.clone(),
                replay_buffer: IO_BUFFER_SIZE,
                key_write_buffer: IO_BUFFER_SIZE,
                value_write_buffer: IO_BUFFER_SIZE,
                block_codec_config,
                max_repair: NZUsize!(10),
                max_pending_acks: NZUsize!(1),
                strategy: self.config.strategy.clone(),
            },
        )
        .await;

        let (fence, _gate) = Fence::new(Epoch::zero());
        let (reshare_actor, reshare_mailbox) = reshare::Actor::new_dkg(
            context.child("reshare"),
            reshare::Config {
                signer: self.config.signer.clone(),
                manager: self.config.manager.clone(),
                blocker: self.config.blocker.clone(),
                participants_provider: StaticParticipants {
                    participants: self.config.participants.clone(),
                    directory: self.config.directory.clone(),
                },
                secret_store: self.config.secret_store,
                strategy: self.config.strategy.clone(),
                registrar: NoopRegistrar(PhantomData),
                marshal: marshal_mailbox.clone(),
                state_sync: StateSyncPlan::disabled(),
                fence,
                namespace: self.config.namespace,
                sharing_mode: self.config.sharing_mode,
                reveal: self.config.reveal,
                mailbox_size: MAILBOX_SIZE,
                partition_prefix: format!("{}-reshare", self.config.partition_prefix),
                max_participants,
                blocks_per_epoch: self.config.blocks_per_epoch,
                batch_verifier: PhantomData::<ed25519::Batch>,
            },
            DkgConfig {
                participants: self.config.participants.clone(),
                directory: self.config.directory.clone(),
                completion: Box::new(move |info| {
                    let _ = completion.send_lossy(Completion { info });
                }),
            },
        );

        let app = reshare::Application::new(
            DkgApp(PhantomData),
            reshare_mailbox.clone(),
            self.config.blocks_per_epoch,
        );
        let deferred = Deferred::new(
            context.child("deferred"),
            app,
            marshal_mailbox.clone(),
            FixedEpocher::new(self.config.blocks_per_epoch),
        );
        let simplex = simplex::Engine::new(
            context.child("simplex"),
            simplex::Config {
                scheme,
                elector: RoundRobin::<Sha256>::default(),
                blocker: self.config.blocker,
                automaton: deferred.clone(),
                relay: deferred,
                reporter: marshal_mailbox.clone(),
                strategy: self.config.strategy,
                partition: format!("{}-simplex", self.config.partition_prefix),
                mailbox_size: MAILBOX_SIZE,
                epoch: Epoch::zero(),
                floor: Floor::Genesis(genesis.digest()),
                replay_buffer: IO_BUFFER_SIZE,
                write_buffer: IO_BUFFER_SIZE,
                page_cache,
                leader_timeout: Duration::from_secs(1),
                certification_timeout: Duration::from_secs(2),
                timeout_retry: Duration::from_millis(500),
                view_retention: ViewDelta::new(10),
                skip: SkipPolicy::Enabled {
                    timeout: Duration::from_secs(5),
                    budget: SkipBudget::Participants,
                },
                fetch_timeout: Duration::from_secs(2),
                forward: ForwardPolicy::Disabled,
                track_historical_votes: false,
            },
        );

        let reshare_handle = reshare_actor.start(dkg);
        let marshal_handle = marshal_actor.start(
            reshare_mailbox,
            buffer_mailbox,
            (backfill_handler, backfill_resolver),
        );
        let simplex_handle = simplex.start(votes, certificates, resolver_network);

        Handle::select([
            buffer_handle,
            reshare_handle,
            marshal_handle,
            simplex_handle,
        ])
        .await
        .expect("failed dkg");
    }
}

#[derive(Clone)]
struct DkgApp<V: Variant, D>(PhantomData<(V, D)>);

impl<E, V, D> Application<E> for DkgApp<V, D>
where
    E: Rng + Spawner + Metrics + Clock,
    V: Variant,
    D: Directory<ed25519::PublicKey>,
{
    type SigningScheme = ConsensusScheme;
    type Context = Context<sha256::Digest, ed25519::PublicKey>;
    type Block = Block<V, D>;
    type Input = reshare::Input<(), V, ed25519::PrivateKey, D>;

    async fn propose(
        &mut self,
        (_, context): (E, Self::Context),
        ancestry: impl Ancestry<Self::Block>,
        input: Self::Input,
    ) -> Option<Self::Block> {
        let parent = ancestry.peek()?.clone();
        let height = parent.height().next();
        Some(Block {
            context,
            parent: parent.digest(),
            height,
            payload: input.payload,
        })
    }

    async fn verify(
        &mut self,
        _: (E, Self::Context),
        _ancestry: impl Ancestry<Self::Block>,
    ) -> bool {
        // The reshare application wrapper validates payload placement and the
        // final block's epoch info before delegating to this stateless leaf.
        true
    }
}

#[derive(Clone)]
struct StaticParticipants<P, D> {
    participants: Set<P>,
    directory: D,
}

impl<P, D> ParticipantsProvider for StaticParticipants<P, D>
where
    P: PublicKey,
    D: Directory<P>,
{
    type PublicKey = P;
    type Directory = D;

    async fn participants(&mut self, _: Epoch) -> Set<Self::PublicKey> {
        self.participants.clone()
    }

    async fn directory(&mut self, _: Epoch, _: Set<Self::PublicKey>) -> Self::Directory {
        self.directory.clone()
    }
}

#[derive(Clone)]
struct NoopRegistrar<V, P>(PhantomData<(V, P)>);

impl<V, P> Registrar for NoopRegistrar<V, P>
where
    V: Variant,
    P: PublicKey,
{
    type Variant = V;
    type PublicKey = P;

    async fn register(&self, _: Epoch, _: SchemeInfo<Self::Variant, Self::PublicKey>) {}
}

fn archive_config<C>(
    prefix: &str,
    name: &str,
    page_cache: CacheRef,
    codec_config: C,
) -> prunable::Config<TwoCap, C> {
    prunable::Config {
        translator: TwoCap,
        metadata_partition: format!("{prefix}-{name}-metadata"),
        key_partition: format!("{prefix}-{name}-key"),
        key_page_cache: page_cache,
        value_partition: format!("{prefix}-{name}-value"),
        compression: None,
        codec_config,
        items_per_section: ARCHIVE_ITEMS_PER_SECTION,
        key_write_buffer: IO_BUFFER_SIZE,
        value_write_buffer: IO_BUFFER_SIZE,
        replay_buffer: IO_BUFFER_SIZE,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use commonware_cryptography::bls12381::primitives::variant::MinPk;

    #[test]
    #[should_panic(expected = "sharing mode must be supported by max supported mode")]
    fn rejects_unsupported_sharing_mode() {
        let config = Config {
            signer: ed25519::PrivateKey::from_seed(0),
            manager: (),
            blocker: (),
            secret_store: (),
            strategy: (),
            namespace: b"test",
            sharing_mode: SharingMode::RootsOfUnity,
            reveal: Reveal::V1,
            max_supported_mode: ModeVersion::v0(),
            partition_prefix: "test".into(),
            participants: Set::default(),
            directory: Unit,
            blocks_per_epoch: NZU64!(1),
        };

        let _ = Engine::<_, MinPk, _, _, _, _, _>::new((), config);
    }
}