//! A `sha3::Sha3_256` Merkle proof example.
//!
//! Run with:
//! ```
//! cargo run --example merkle-proof
//! ```
//!
//! Check the command line options with:
//! ```
//! cargo run --example merkle-proof -- -h
//! ```

use std::collections::HashSet;
use std::iter;
use std::num::NonZeroUsize;
use std::sync::mpsc::{channel, Sender};
use std::thread;

use clap::Parser;
use digest::generic_array::{ArrayLength, GenericArray};
use digest::{Digest, OutputSizeUser};
use merkle_lite::{MerkleProof, MerkleTree};
use rand_core::RngCore;
use tracing::{info, instrument, trace};

#[derive(Debug, Parser)]
#[command(author, version, about, long_about = None)]
struct Args {
    /// Number of leaves.
    #[arg(short = 'n', long, default_value_t = 10_000)]
    leaves: usize,

    /// Depth of the tree, instead of the number of leaves.
    #[arg(short = 'd', long)]
    depth: Option<NonZeroUsize>,

    /// Number of random leaves to check for the proof of inclusion.
    #[arg(short = 'p', long, default_value_t = 5)]
    proof_of_inclusion: usize,
}

fn main() {
    tracing_subscriber::fmt::init();

    // Gets a number of leaves.
    let args = Args::parse();
    let leaf_len = args
        .depth
        .and_then(|depth| 1usize.checked_shl(depth.get() as u32 - 1))
        .unwrap_or(args.leaves);

    info!("Merkle proof of inclusion");

    // Random indices for the proof of inclusion check below.
    let leaf_indices: HashSet<_> = iter::repeat_with(|| fastrand::usize(..leaf_len))
        .take(args.proof_of_inclusion)
        .collect();

    // A channel to receive the Merkle proof generated by `proof_server`.
    let (tx, rx) = channel();

    // Spawns the Merkle proof server.
    thread::spawn(move || server::<sha3::Sha3_256>(tx, leaf_len, leaf_indices));

    // Waits for the Merkle proof from the server.
    let (proof, root, leaves) = rx.recv().expect("server crashed");

    // Verifies the proof of inclusion.
    assert_eq!(
        proof.verify(&leaves).expect("verify failed").as_ref(),
        root.as_slice(),
    );
    info!("Merkle proof was successfully verified");
}

/// A Merkle root and leaf hashes for the proof of inclusion.
type TreeRoot = Vec<u8>;
type TreeLeaves = Vec<(usize, Vec<u8>)>;

/// A Merkle proof server.
#[instrument(name = "merkle_proof_server", skip(tx))]
fn server<B>(
    tx: Sender<(MerkleProof<B>, TreeRoot, TreeLeaves)>,
    leaf_length: usize,
    proof_of_leaf_indices: HashSet<usize>,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>>
where
    B: Digest + 'static,
    <<B as OutputSizeUser>::OutputSize as ArrayLength<u8>>::ArrayType: Copy,
{
    // Composes a Merkle tree for `leaf_length` random leaves.
    info!("Server started");
    let tree: MerkleTree<B> =
        iter::repeat(GenericArray::<u8, <B as OutputSizeUser>::OutputSize>::default())
            .map(|mut hash| {
                rand_core::OsRng.fill_bytes(&mut hash);
                hash
            })
            .take(leaf_length)
            .collect();
    trace!(?tree, "Got the tree");

    let root = tree.root().ok_or("merkle root")?.to_vec();
    trace!(?root, "Got the root");

    // Gets the Merkle root and proof.
    let tree_leaves = tree.get_leaves();
    let proof = tree.proof(&proof_of_leaf_indices).ok_or("merkle proof")?;
    trace!(?proof, "Got the proof");

    // Gets the leaf index and the hash requested by the caller.
    let leaves: TreeLeaves = proof_of_leaf_indices
        .iter()
        .map(|index| (*index, tree_leaves[*index].as_slice().to_vec()))
        .collect();
    trace!(leaves_len = leaves.len(), "Got the leaves");

    // Returns a result over the channel.
    tx.send((proof, root, leaves))?;
    info!("Server Finished");

    Ok(())
}