//! `bootstrap` subcommand: one-shot glue DKG for the epoch-0 committee.

use crate::{
    config::{NetworkConfig, NodeConfig},
    types::{
        self, BACKFILL_CHANNEL, BLOCKS_PER_EPOCH, BROADCAST_CHANNEL, CERTIFICATE_CHANNEL,
        DKG_CHANNEL, IO_BUFFER_SIZE, ITEMS_PER_SECTION, MAILBOX_SIZE, MAX_SUPPORTED_MODE,
        MESSAGE_RATE, NAMESPACE, PAGE_CACHE_SIZE, PAGE_SIZE, Participants, Partition,
        RESOLVER_CHANNEL, REVEAL, SHARING_MODE, Secrets, VOTE_CHANNEL,
    },
    validator,
};
use clap::Args;
use commonware_consensus::types::Epoch;
use commonware_cryptography::{
    ChaCha20Poly1305, bls12381::primitives::variant::MinSig, ed25519::PublicKey,
};
use commonware_glue::dkg::{
    SecretStore as _, bootstrap,
    types::{EpochInfo, EpochOutcome},
};
use commonware_p2p::{
    authenticated::{self, discovery},
    max_message_size,
};
use commonware_runtime::{Handle, Strategizer, Supervisor as _, buffer::paged::CacheRef, tokio};
use commonware_stream::{
    cups::{self, Cups},
    sake::{self, Sake},
};
use commonware_utils::{NZDuration, NZU32, NZUsize, sequence::Unit};
use std::{
    fs,
    path::{Path, PathBuf},
    time::Duration,
};
use tracing::{error, info};

type ReshareEpochInfo = EpochInfo<MinSig, PublicKey>;

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

/// Run the one-shot DKG, write the resulting genesis, and keep serving until
/// stopped.
#[derive(Args)]
pub struct Bootstrap {
    /// Validator node directory containing config and runtime storage.
    #[arg(long, default_value = "./data/validator-0")]
    pub node_dir: PathBuf,
}

/// Run the [`bootstrap::Engine`], distribute the genesis artifact on completion,
/// and keep serving peers that have not completed.
pub async fn run(context: tokio::Context, args: Bootstrap) {
    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 participants = Participants::new(&network).expect("invalid participants");
    let local = node.public_key();
    let bootstrappers = network.bootstrappers(&local);
    let max_peers_per_set = authenticated::peer_set_limit(&network.participants, &local);
    let committee = participants.get(Epoch::zero());
    let committee_size =
        NZU32!(u32::try_from(committee.len()).expect("committee size must fit in u32"));
    let payload = bootstrap::max_payload_size::<MinSig, Unit>(committee_size, MAX_SUPPORTED_MODE);

    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,
        max_message_size(&[payload]),
        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 = p2p.register(VOTE_CHANNEL, MESSAGE_RATE);
    let certificate = p2p.register(CERTIFICATE_CHANNEL, MESSAGE_RATE);
    let resolver = p2p.register(RESOLVER_CHANNEL, MESSAGE_RATE);
    let backfill = p2p.register(BACKFILL_CHANNEL, MESSAGE_RATE);
    let broadcast = p2p.register(BROADCAST_CHANNEL, MESSAGE_RATE);
    let dkg = p2p.register(DKG_CHANNEL, MESSAGE_RATE);

    let strategy = context.strategy(NZUsize!(2));
    let mut store = Secrets::init(context.child("secrets"), PARTITION).await;
    let engine = bootstrap::Engine::new(
        context.child("bootstrap"),
        bootstrap::Config {
            signer: node.signer,
            manager: oracle.clone(),
            blocker: oracle.clone(),
            secret_store: store.clone(),
            strategy,
            namespace: NAMESPACE,
            sharing_mode: SHARING_MODE,
            reveal: REVEAL,
            max_supported_mode: MAX_SUPPORTED_MODE,
            partition_prefix: "bootstrap".to_string(),
            page_cache: CacheRef::new(PAGE_SIZE, PAGE_CACHE_SIZE),
            write_buffer: IO_BUFFER_SIZE,
            replay_buffer: IO_BUFFER_SIZE,
            mailbox_size: MAILBOX_SIZE,
            fetch_timeout: NZDuration!(Duration::from_secs(2)),
            fetch_retry_delay: Duration::from_millis(100),
            items_per_section: ITEMS_PER_SECTION,
            participants: committee,
            directory: Unit,
            blocks_per_epoch: BLOCKS_PER_EPOCH,
        },
    );

    let p2p_handle = p2p.start();
    let (engine_handle, completion) =
        engine.start(vote, certificate, resolver, backfill, broadcast, dkg);
    let info = completion
        .await
        .expect("bootstrap completion dropped")
        .info
        .expect("bootstrap DKG failed");

    // Hand only the epoch-0 share to the validator store. The bootstrap store
    // also holds this ceremony's epoch-0 seed and dealings, which the
    // validator's first reshare must not reuse.
    if let Some(share) = store.get_share(Epoch::zero()).await {
        let mut handoff = Secrets::init(context.child("handoff"), validator::PARTITION).await;
        handoff.put_share(Epoch::zero(), share).await;
    }

    let mut genesis = info;
    genesis.outcome = EpochOutcome::Success;
    genesis.next_players = participants.get(genesis.epoch.next());
    let written = write_genesis_to_sibling_validators(&args.node_dir, &local, &network, &genesis)
        .expect("failed to write genesis");
    info!(
        epoch = genesis.epoch.get(),
        players = genesis.players.len(),
        next_players = genesis.next_players.len(),
        directories = written,
        "bootstrap complete, serving peers until stopped"
    );

    // Keep serving the one-shot chain so participants that have not completed
    // can catch up.
    if let Err(err) = Handle::select([p2p_handle, engine_handle]).await {
        error!(?err, "bootstrap task failed");
    }
}

/// Write `genesis` into every sibling validator directory that belongs to
/// `network`, except those of players other than `local`, or into `node_dir`
/// alone when none are found.
fn write_genesis_to_sibling_validators(
    node_dir: &Path,
    local: &PublicKey,
    network: &NetworkConfig,
    genesis: &ReshareEpochInfo,
) -> anyhow::Result<usize> {
    let Some(root) = node_dir.parent() else {
        types::write_genesis(node_dir, genesis)?;
        return Ok(1);
    };

    let mut written = 0;
    for entry in fs::read_dir(root)? {
        let candidate = entry?.path();
        if !candidate.is_dir() {
            continue;
        }
        let Ok(node) = NodeConfig::load(&candidate) else {
            continue;
        };
        let public_key = node.public_key();
        if !network.participants.contains(&public_key) {
            continue;
        }
        if public_key != *local && genesis.players.position(&public_key).is_some() {
            continue;
        }
        types::write_genesis(&candidate, genesis)?;
        written += 1;
    }

    if written == 0 {
        types::write_genesis(node_dir, genesis)?;
        return Ok(1);
    }
    Ok(written)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::{self, PeerConfig};
    use commonware_cryptography::{
        Signer as _, bls12381::dkg::feldman_desmedt::deal, ed25519::PrivateKey,
    };
    use commonware_math::algebra::Random;
    use commonware_utils::{N3f1, ordered::Set, test_rng};

    #[test]
    fn writes_bootstrap_genesis_to_own_and_non_player_validators() {
        let root =
            std::env::temp_dir().join(format!("commonware-dkg-bootstrap-{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&root);
        std::fs::create_dir_all(&root).unwrap();

        let mut rng = test_rng();
        let signers = (0..4)
            .map(|_| PrivateKey::random(&mut rng))
            .collect::<Vec<_>>();
        let network = NetworkConfig {
            participants: signers.iter().map(|signer| signer.public_key()).collect(),
            committee_size: 2,
            peers: signers
                .iter()
                .enumerate()
                .map(|(i, signer)| PeerConfig {
                    public_key: signer.public_key(),
                    dial: format!("127.0.0.1:{}", 4300 + i)
                        .parse()
                        .expect("valid address"),
                })
                .collect(),
        };

        for (i, signer) in signers.into_iter().enumerate() {
            let node_dir = root.join(format!("validator-{i}"));
            std::fs::create_dir_all(&node_dir).unwrap();
            config::write_json(
                &node_dir.join("node.json"),
                &NodeConfig {
                    signer,
                    listen: format!("127.0.0.1:{}", 4300 + i)
                        .parse()
                        .expect("valid address"),
                    dial: format!("127.0.0.1:{}", 4300 + i)
                        .parse()
                        .expect("valid address"),
                },
            )
            .unwrap();
            config::write_json(&node_dir.join("network.json"), &network).unwrap();
        }

        let players = Set::from_iter_dedup(network.participants.iter().take(2).cloned());
        let (output, _shares) =
            deal::<MinSig, _, N3f1>(test_rng(), SHARING_MODE, players.clone()).unwrap();
        let genesis = EpochInfo {
            outcome: EpochOutcome::Success,
            epoch: Epoch::zero(),
            output,
            players,
            next_players: Set::from_iter_dedup(
                network.participants.iter().skip(1).take(2).cloned(),
            ),
            directory: Unit,
        };

        let written = write_genesis_to_sibling_validators(
            &root.join("validator-0"),
            &network.participants[0],
            &network,
            &genesis,
        )
        .unwrap();

        // Player 0 writes its own directory and every non-player's, but not
        // the directory of player 1.
        assert_eq!(written, 3);
        assert!(!types::genesis_path(&root.join("validator-1")).exists());
        for i in [0, 2, 3] {
            assert_eq!(
                types::read_genesis(&root.join(format!("validator-{i}"))).unwrap(),
                genesis
            );
        }
        let _ = std::fs::remove_dir_all(root);
    }
}