//! `validator` subcommand: run a validator node.

use crate::{
    application::App,
    bootstrap,
    config::{NetworkConfig, NodeConfig},
    types::{
        self, BACKFILL_CHANNEL, BLOCKS_PER_EPOCH, BROADCAST_CHANNEL, Block, CERTIFICATE_CHANNEL,
        DKG_CHANNEL, DKG_PROBE_CHANNEL, DynamicProvider, IO_BUFFER_SIZE, ITEMS_PER_SECTION,
        LogReporter, MAILBOX_SIZE, MAX_PARTICIPANTS, MAX_SUPPORTED_MODE, MESSAGE_RATE, NAMESPACE,
        PAGE_CACHE_SIZE, PAGE_SIZE, Participants, Partition, QMDB_CHANNEL, RESOLVER_CHANNEL,
        REVEAL, Registrar, SHARING_MODE, Scheme, Secrets, VOTE_CHANNEL,
    },
};
use clap::Args;
use commonware_broadcast::buffered;
use commonware_consensus::{
    Reporters,
    marshal::{
        self,
        core::Actor as MarshalActor,
        resolver::p2p as marshal_resolver,
        standard::{Deferred, Standard},
    },
    simplex::{
        SkipBudget,
        config::{ForwardPolicy, SkipPolicy},
        elector::RoundRobin,
    },
    types::{Epoch, FixedEpocher, ViewDelta},
};
use commonware_cryptography::{
    ChaCha20Poly1305, Digested, bls12381::primitives::variant::MinSig, ed25519, sha256::Sha256,
};
use commonware_glue::{
    dkg::{
        SecretStore as _,
        fence::Fence,
        orchestrator, probe, reshare,
        state_sync::{Config as StateSyncConfig, Plan as StateSyncPlan, StateSync},
        types::{EpochInfo, SchemeInfo},
    },
    stateful::{
        Config as StatefulConfig, Stateful, SyncPlan,
        db::{DatabaseSet, p2p as qmdb_resolver},
    },
};
use commonware_macros::boxed;
use commonware_p2p::authenticated::{
    self,
    discovery::{self, Oracle},
};
use commonware_parallel::Sequential;
use commonware_runtime::{Handle, Spawner, Supervisor as _, buffer::paged::CacheRef, tokio};
use commonware_storage::{archive::prunable, translator::TwoCap};
use commonware_stream::{
    cups::{self, Cups},
    sake::{self, Sake},
};
use commonware_utils::{NZDuration, NZUsize, sequence::Unit};
use rand_core::CryptoRng;
use std::{path::PathBuf, time::Duration};
use tracing::error;

/// Start a validator node.
#[derive(Args)]
pub struct Validator {
    /// Validator node directory containing config, genesis, and runtime storage.
    #[arg(long, default_value = "./data/validator-0")]
    pub node_dir: PathBuf,

    /// Run one-time peer state sync for a new late joiner.
    #[arg(long, default_value_t = false)]
    pub state_sync: bool,
}

/// Partition of the validator's [`Secrets`] store.
pub const PARTITION: Partition = Partition::Validator;

/// Start every validator actor and run until one stops.
#[boxed]
pub async fn run(context: tokio::Context, args: Validator) {
    let node = NodeConfig::load(&args.node_dir).expect("failed to load node config");
    let network = NetworkConfig::load(&args.node_dir).expect("failed to load network config");
    network.validate().expect("invalid network config");
    let genesis_info = types::read_genesis(&args.node_dir).expect("genesis is required");

    // A player's genesis comes only from its own `bootstrap`, which hands the
    // epoch-0 share to `secrets` before writing genesis. Nothing in the
    // bootstrap store is needed after that, so erase it.
    Secrets::init(context.child("bootstrap"), bootstrap::PARTITION)
        .await
        .destroy()
        .await;
    let participants = Participants::new(&network).expect("invalid participants");
    let local = node.public_key();
    let partition_prefix = "validator";
    let page_cache = CacheRef::new(PAGE_SIZE, PAGE_CACHE_SIZE);
    let bootstrappers = network.bootstrappers(&local);
    let max_peers_per_set = authenticated::peer_set_limit(&network.participants, &local);

    let p2p_config = discovery::Config::local(
        Cups::<_, ChaCha20Poly1305>::new(
            Sake {
                signer: node.signer.clone(),
                synchrony_bound: Duration::from_secs(5),
                max_handshake_age: Duration::from_secs(10),
                version: sake::Version::V1,
            },
            cups::Version::V1,
        ),
        &[NAMESPACE, b"_P2P"].concat(),
        node.listen,
        node.dial,
        bootstrappers,
        max_peers_per_set,
        types::max_message_size(),
        MAILBOX_SIZE,
    );

    let (mut p2p, oracle) = discovery::Network::new(context.child("network"), p2p_config);

    // Channel rates are enforced independently per peer. The network derives each shared inbound
    // mailbox capacity from the retained-peer bound and quota burst size.
    let vote_network = p2p.register(VOTE_CHANNEL, MESSAGE_RATE);
    let certificate_network = p2p.register(CERTIFICATE_CHANNEL, MESSAGE_RATE);
    let resolver_network = p2p.register(RESOLVER_CHANNEL, MESSAGE_RATE);
    let backfill_network = p2p.register(BACKFILL_CHANNEL, MESSAGE_RATE);
    let broadcast_network = p2p.register(BROADCAST_CHANNEL, MESSAGE_RATE);
    let qmdb_network = p2p.register(QMDB_CHANNEL, MESSAGE_RATE);
    let dkg_network = p2p.register(DKG_CHANNEL, MESSAGE_RATE);
    let dkg_probe_network = p2p.register(DKG_PROBE_CHANNEL, MESSAGE_RATE);
    let p2p_handle = p2p.start();

    let provider = DynamicProvider::default();
    let mut store = Secrets::init(context.child("secrets"), PARTITION).await;
    let share = store.get_share(Epoch::zero()).await;
    provider.register(
        Epoch::zero(),
        types::scheme(SchemeInfo::new(
            &local,
            genesis_info.output.players().clone(),
            genesis_info.output.public().clone(),
            share,
        )),
    );

    let resolver = marshal_resolver::init(
        context.child("marshal_resolver"),
        marshal_resolver::Config {
            public_key: local.clone(),
            peer_provider: oracle.clone(),
            blocker: oracle.clone(),
            mailbox_size: MAILBOX_SIZE,
            fetch_timeout: NZDuration!(Duration::from_secs(2)),
            fetch_retry_delay: Duration::from_millis(100),
            priority_requests: false,
            priority_responses: false,
        },
        backfill_network,
    );

    let (broadcast_engine, buffer) = buffered::Engine::new(
        context.child("broadcast"),
        buffered::Config {
            public_key: local.clone(),
            mailbox_size: MAILBOX_SIZE,
            deque_size: 16,
            priority: false,
            codec_config: (),
            peer_provider: oracle.clone(),
            strategy: Sequential,
        },
    );
    let broadcast_handle = broadcast_engine.start(broadcast_network);

    let finalizations_by_height = prunable::Archive::init(
        context.child("finalizations_by_height"),
        archive_config(partition_prefix, "finalizations", page_cache.clone(), ()),
    )
    .await
    .expect("finalizations archive");
    let finalized_blocks = prunable::Archive::init(
        context.child("finalized_blocks"),
        archive_config(partition_prefix, "blocks", page_cache.clone(), ()),
    )
    .await
    .expect("blocks archive");

    let genesis_target =
        <types::Database<tokio::Context> as DatabaseSet<tokio::Context>>::initial_sync_targets();
    let genesis = Digested::compute(
        Block::genesis(
            network.participants[0].clone(),
            genesis_info.clone(),
            genesis_target,
        ),
        &Sequential,
    )
    .await;
    let stateful_startup = context.child("stateful_startup");
    let mut plan = SyncPlan::init(stateful_startup.child("plan"), partition_prefix).await;
    let (probe_actor, probe_mailbox) = probe::Actor::new(probe_config(
        context.child("dkg_probe"),
        oracle.clone(),
        &genesis_info,
        &plan,
    ));
    let probe_handle = probe_actor.start(dkg_probe_network);

    let should_state_sync = plan.should_sync(args.state_sync);
    let probe_artifact = if should_state_sync {
        let artifact = probe_mailbox.subscribe().await.expect("probe stopped");
        provider.register(
            artifact.info.epoch,
            types::scheme(SchemeInfo::new(
                &local,
                artifact.info.output.players().clone(),
                artifact.info.output.public().clone(),
                None,
            )),
        );
        plan = plan.set_floor(artifact.floor.clone()).await;
        Some(artifact)
    } else {
        None
    };

    let (marshal_actor, marshal, floor) = MarshalActor::init(
        context.child("marshal"),
        finalizations_by_height,
        finalized_blocks,
        marshal::Config {
            provider: provider.clone(),
            epocher: FixedEpocher::new(BLOCKS_PER_EPOCH),
            start: plan.marshal_start(genesis.clone()),
            partition_prefix: partition_prefix.to_string(),
            mailbox_size: MAILBOX_SIZE,
            view_retention: ViewDelta::new(10),
            prunable_items_per_section: ITEMS_PER_SECTION,
            page_cache: page_cache.clone(),
            replay_buffer: types::IO_BUFFER_SIZE,
            key_write_buffer: types::IO_BUFFER_SIZE,
            value_write_buffer: types::IO_BUFFER_SIZE,
            block_codec_config: (),
            max_repair: NZUsize!(10),
            max_pending_acks: NZUsize!(1),
            strategy: Sequential,
        },
    )
    .await;

    let (qmdb_actor, qmdb_sync_resolver) = qmdb_resolver::Actor::new(
        context.child("qmdb_resolver"),
        qmdb_resolver::Config {
            peer_provider: oracle.clone(),
            blocker: oracle.clone(),
            database: None,
            mailbox_size: MAILBOX_SIZE,
            me: Some(local.clone()),
            fetch_timeout: NZDuration!(Duration::from_secs(2)),
            fetch_retry_delay: Duration::from_millis(100),
            max_serve_ops: types::SYNC_BATCH_SIZE,
            priority_requests: false,
            priority_responses: false,
        },
    );
    let qmdb_handle = qmdb_actor.start(qmdb_network);

    let state_sync = probe_artifact.map(|artifact| {
        let floor = plan
            .floor()
            .cloned()
            .expect("state sync startup must have floor");
        StateSync {
            info: artifact.info,
            floor,
        }
    });
    let state_sync = StateSyncPlan::init(
        context.child("dkg_state_sync_plan"),
        StateSyncConfig {
            partition_prefix: partition_prefix.to_string(),
            max_participants: MAX_PARTICIPANTS,
            max_supported_mode: MAX_SUPPORTED_MODE,
        },
        state_sync,
    )
    .await;

    // The orchestrator enters an epoch only once the reshare actor registers
    // the epoch's scheme and marks the fence, so it never starts an engine from
    // the state-sync verifier registered above for marshal.
    let (fence, gate) = Fence::new();
    let (reshare_actor, reshare_mailbox) = reshare::Actor::new(
        context.child("reshare"),
        reshare::Config {
            signer: node.signer,
            manager: oracle.clone(),
            blocker: oracle.clone(),
            participants_provider: participants,
            secret_store: store,
            strategy: Sequential,
            registrar: Registrar::new(provider.clone()),
            marshal: marshal.clone(),
            state_sync: state_sync.clone(),
            fence,
            namespace: NAMESPACE,
            sharing_mode: SHARING_MODE,
            reveal: REVEAL,
            mailbox_size: MAILBOX_SIZE,
            partition_prefix: format!("{partition_prefix}-reshare"),
            page_cache: page_cache.clone(),
            write_buffer: IO_BUFFER_SIZE,
            replay_buffer: IO_BUFFER_SIZE,
            max_participants: MAX_PARTICIPANTS,
            blocks_per_epoch: BLOCKS_PER_EPOCH,
        },
    );
    let reshare_handle = reshare_actor.start(dkg_network);

    let (stateful_actor, stateful_mailbox) = Stateful::new(
        context.child("stateful"),
        StatefulConfig {
            application: App::new(genesis.inner().clone()),
            db_config: types::db_config(partition_prefix, page_cache.clone()),
            provider: (),
            marshal: (marshal.clone(), floor),
            mailbox_size: MAILBOX_SIZE,
            plan,
            resolvers: qmdb_sync_resolver,
            sync_config: types::sync_config(),
            prune_config: None,
            strategy: Sequential,
        },
    );

    // The reshare wrapper drives the payload for the stateful application.
    let deferred = Deferred::new(
        context.child("deferred"),
        reshare::Application::new(
            stateful_mailbox.clone(),
            reshare_mailbox.clone(),
            BLOCKS_PER_EPOCH,
        ),
        marshal.clone(),
        FixedEpocher::new(BLOCKS_PER_EPOCH),
    );
    let (orchestrator_actor, orchestrator_mailbox) = orchestrator::Actor::new(
        context.child("orchestrator"),
        orchestrator::Config {
            oracle: oracle.clone(),
            manager: oracle.clone(),
            provider: provider.clone(),
            marshal: marshal.clone(),
            application: deferred,
            strategy: Sequential,
            simplex: orchestrator::SimplexConfig {
                elector: RoundRobin::<Sha256>::default(),
                mailbox_size: NZUsize!(3),
                replay_buffer: IO_BUFFER_SIZE,
                write_buffer: IO_BUFFER_SIZE,
                page_cache: page_cache.clone(),
                leader_timeout: Duration::from_secs(1),
                certification_timeout: Duration::from_secs(2),
                timeout_retry: Duration::from_millis(500),
                fetch_timeout: NZDuration!(Duration::from_secs(2)),
                fetch_retry_delay: Duration::from_millis(100),
                view_retention: ViewDelta::new(10),
                skip: SkipPolicy::Enabled {
                    timeout: Duration::from_secs(5),
                    budget: SkipBudget::Participants,
                },
                forward: ForwardPolicy::Disabled,
                track_historical_votes: false,
            },
            gate,
            state_sync,
            blocks_per_epoch: BLOCKS_PER_EPOCH,
            mailbox_size: MAILBOX_SIZE,
            partition_prefix: format!("{partition_prefix}-orchestrator"),
        },
    );
    let orchestrator_handle =
        orchestrator_actor.start(vote_network, certificate_network, resolver_network);

    let reporters = Reporters::from((
        stateful_mailbox.clone(),
        Reporters::from((
            orchestrator_mailbox,
            Reporters::from((reshare_mailbox, LogReporter)),
        )),
    ));
    let marshal_handle = marshal_actor.start(reporters, buffer, resolver);
    probe_mailbox.attach(marshal.clone());
    let stateful_handle = stateful_actor.start();

    if let Err(err) = Handle::select([
        p2p_handle,
        broadcast_handle,
        probe_handle,
        qmdb_handle,
        reshare_handle,
        orchestrator_handle,
        marshal_handle,
        stateful_handle,
    ])
    .await
    {
        error!(?err, "validator task failed");
    }
}

/// Configure the DKG probe with `plan`'s persisted state sync floor.
fn probe_config<E>(
    context: E,
    oracle: Oracle<ed25519::PublicKey>,
    genesis: &EpochInfo<MinSig, ed25519::PublicKey>,
    plan: &SyncPlan<E, Scheme, Standard<Block>>,
) -> probe::Config<
    E,
    Oracle<ed25519::PublicKey>,
    Scheme,
    Standard<Block>,
    Sequential,
    Oracle<ed25519::PublicKey>,
>
where
    E: Spawner + CryptoRng + commonware_storage::Context,
{
    probe::Config {
        context,
        manager: oracle.clone(),
        bootstrap: probe::Bootstrap {
            epoch: Epoch::zero(),
            participants: genesis.participants(),
            directory: Unit,
        },
        floor: plan.floor().cloned(),
        verifier: Scheme::certificate_verifier(NAMESPACE, *genesis.output.public().public()),
        genesis: genesis.clone(),
        strategy: Sequential,
        blocker: oracle,
        blocks_per_epoch: BLOCKS_PER_EPOCH,
        retry_timeout: NZDuration!(Duration::from_millis(500)),
        mailbox_size: MAILBOX_SIZE,
        block_codec_config: (),
    }
}

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: 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_consensus::{
        simplex::types::{Finalization, Finalize, Proposal},
        types::{Round, View},
    };
    use commonware_cryptography::{
        ChaCha20Poly1305, Hasher as _, Signer as _, bls12381::dkg::feldman_desmedt::deal,
    };
    use commonware_glue::dkg::types::EpochOutcome;
    use commonware_runtime::{Runner as _, deterministic};
    use commonware_utils::{N3f1, TestRng, non_empty, ordered::Set};
    use futures::{FutureExt as _, future::pending};
    use std::{
        net::SocketAddr,
        sync::{
            Arc,
            atomic::{AtomicUsize, Ordering},
        },
    };

    struct CountDrop(Arc<AtomicUsize>);

    impl Drop for CountDrop {
        fn drop(&mut self) {
            self.0.fetch_add(1, Ordering::Relaxed);
        }
    }

    fn pending_handle(dropped: Arc<AtomicUsize>) -> Handle<()> {
        let count_drop = CountDrop(dropped);
        Handle::from_future(async move {
            let _count_drop = count_drop;
            pending().await
        })
    }

    #[test]
    fn successful_actor_completion_stops_validator() {
        let dropped = Arc::new(AtomicUsize::new(0));

        // Model a clean actor exit alongside siblings that would otherwise run forever.
        let actors = [
            Handle::ready(Ok(())),
            pending_handle(dropped.clone()),
            pending_handle(dropped.clone()),
            pending_handle(dropped.clone()),
            pending_handle(dropped.clone()),
            pending_handle(dropped.clone()),
            pending_handle(dropped.clone()),
            pending_handle(dropped.clone()),
        ];

        // Supervision must complete and abort every pending sibling.
        assert!(matches!(
            Handle::select(actors).now_or_never(),
            Some(Ok(()))
        ));
        assert_eq!(dropped.load(Ordering::Relaxed), 7);
    }

    /// A node restarting during state sync configures its probe with the floor
    /// the interrupted sync persisted.
    #[test]
    fn probe_resumes_persisted_floor() {
        deterministic::Runner::default().start(|context| async move {
            // Deal a genesis committee and sign a floor in epoch 1 with it.
            let signers = (0..4)
                .map(ed25519::PrivateKey::from_seed)
                .collect::<Vec<_>>();
            let players = Set::from_iter_dedup(signers.iter().map(|signer| signer.public_key()));
            let (output, shares) =
                deal::<MinSig, _, N3f1>(TestRng::new(0), SHARING_MODE, players.clone())
                    .expect("genesis deal");
            let schemes = shares
                .values()
                .iter()
                .map(|share| {
                    Scheme::signer(
                        NAMESPACE,
                        players.clone(),
                        output.public().clone(),
                        share.clone(),
                    )
                    .expect("genesis signer share")
                })
                .collect::<Vec<_>>();
            let proposal = Proposal::new(
                Round::new(Epoch::new(1), View::new(1)),
                View::zero(),
                Sha256::hash(&[b"floor"]),
            );
            let finalizes = schemes
                .iter()
                .map(|scheme| Finalize::sign(scheme, proposal.clone()).expect("sign finalize"))
                .collect::<Vec<_>>();
            let floor = Finalization::from_finalizes(
                &schemes[0],
                non_empty![@finalizes.iter()],
                &Sequential,
            )
            .expect("finalization quorum");
            let genesis = EpochInfo {
                outcome: EpochOutcome::Success,
                epoch: Epoch::zero(),
                output,
                players: players.clone(),
                next_players: players,
                directory: Unit,
            };

            // Persist the floor, then reload the plan as a restarted node.
            let plan = SyncPlan::<_, Scheme, Standard<Block>>::init(context.child("plan"), "probe")
                .await
                .set_floor(floor.clone())
                .await;
            drop(plan);
            let plan = SyncPlan::init(context.child("restart"), "probe").await;

            // The probe carries the persisted floor.
            let address = SocketAddr::from(([127, 0, 0, 1], 3000));
            let p2p_config = discovery::Config::local(
                Cups::<_, ChaCha20Poly1305>::new(
                    Sake {
                        signer: signers[0].clone(),
                        synchrony_bound: Duration::from_secs(5),
                        max_handshake_age: Duration::from_secs(10),
                        version: sake::Version::V1,
                    },
                    cups::Version::V1,
                ),
                NAMESPACE,
                address,
                address,
                Vec::new(),
                NZUsize!(signers.len()),
                types::max_message_size(),
                MAILBOX_SIZE,
            );
            let (_, oracle) = discovery::Network::new(context.child("network"), p2p_config);
            let config = probe_config(context.child("probe"), oracle, &genesis, &plan);
            assert_eq!(config.floor, Some(floor));
        });
    }
}