// This file is part of Substrate.

// Copyright (C) Parity Technologies (UK) Ltd.
// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0

// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

//! BABE authority selection and slot claiming.

use super::{Epoch, AUTHORING_SCORE_LENGTH, AUTHORING_SCORE_VRF_CONTEXT};
use codec::Encode;
use sc_consensus_epochs::Epoch as EpochT;
use sp_application_crypto::AppCrypto;
use sp_consensus_babe::{
	digests::{PreDigest, PrimaryPreDigest, SecondaryPlainPreDigest, SecondaryVRFPreDigest},
	make_vrf_sign_data, AuthorityId, BabeAuthorityWeight, Randomness, Slot,
};
use sp_core::{
	crypto::{ByteArray, Wraps},
	U256,
};
use sp_keystore::KeystorePtr;

/// Calculates the primary selection threshold for a given authority, taking
/// into account `c` (`1 - c` represents the probability of a slot being empty).
pub(super) fn calculate_primary_threshold(
	c: (u64, u64),
	authorities: &[(AuthorityId, BabeAuthorityWeight)],
	authority_index: usize,
) -> u128 {
	use num_bigint::BigUint;
	use num_rational::BigRational;
	use num_traits::{cast::ToPrimitive, identities::One};

	// Prevent div by zero and out of bounds access.
	// While Babe's pallet implementation that ships with FRAME performs a sanity check over
	// configuration parameters, this is not sufficient to guarantee that `c.1` is non-zero
	// (i.e. third party implementations are possible).
	if c.1 == 0 || authority_index >= authorities.len() {
		return 0
	}

	let c = c.0 as f64 / c.1 as f64;

	let theta = authorities[authority_index].1 as f64 /
		authorities.iter().map(|(_, weight)| weight).sum::<u64>() as f64;

	assert!(theta > 0.0, "authority with weight 0.");

	// NOTE: in the equation `p = 1 - (1 - c)^theta` the value of `p` is always
	// capped by `c`. For all practical purposes `c` should always be set to a
	// value < 0.5, as such in the computations below we should never be near
	// edge cases like `0.999999`.

	let p = BigRational::from_float(1f64 - (1f64 - c).powf(theta)).expect(
		"returns None when the given value is not finite; \
		 c is a configuration parameter defined in (0, 1]; \
		 theta must be > 0 if the given authority's weight is > 0; \
		 theta represents the validator's relative weight defined in (0, 1]; \
		 powf will always return values in (0, 1] given both the \
		 base and exponent are in that domain; \
		 qed.",
	);

	let numer = p.numer().to_biguint().expect(
		"returns None when the given value is negative; \
		 p is defined as `1 - n` where n is defined in (0, 1]; \
		 p must be a value in [0, 1); \
		 qed.",
	);

	let denom = p.denom().to_biguint().expect(
		"returns None when the given value is negative; \
		 p is defined as `1 - n` where n is defined in (0, 1]; \
		 p must be a value in [0, 1); \
		 qed.",
	);

	((BigUint::one() << 128usize) * numer / denom).to_u128().expect(
		"returns None if the underlying value cannot be represented with 128 bits; \
		 we start with 2^128 which is one more than can be represented with 128 bits; \
		 we multiple by p which is defined in [0, 1); \
		 the result must be lower than 2^128 by at least one and thus representable with 128 bits; \
		 qed.",
	)
}

/// Get the expected secondary author for the given slot and with given
/// authorities. This should always assign the slot to some authority unless the
/// authorities list is empty.
pub(super) fn secondary_slot_author(
	slot: Slot,
	authorities: &[(AuthorityId, BabeAuthorityWeight)],
	randomness: Randomness,
) -> Option<&AuthorityId> {
	if authorities.is_empty() {
		return None
	}

	let rand =
		U256::from_big_endian(&(randomness, slot).using_encoded(sp_crypto_hashing::blake2_256));

	let authorities_len = U256::from(authorities.len());
	let idx = rand % authorities_len;

	let expected_author = authorities.get(idx.as_u32() as usize).expect(
		"authorities not empty; index constrained to list length; \
				this is a valid index; qed",
	);

	Some(&expected_author.0)
}

/// Claim a secondary slot if it is our turn to propose, returning the
/// pre-digest to use when authoring the block, or `None` if it is not our turn
/// to propose.
fn claim_secondary_slot(
	slot: Slot,
	epoch: &Epoch,
	keys: &[(AuthorityId, usize)],
	keystore: &KeystorePtr,
	author_secondary_vrf: bool,
) -> Option<(PreDigest, AuthorityId)> {
	if epoch.authorities.is_empty() {
		return None
	}

	let mut epoch_index = epoch.epoch_index;
	if epoch.end_slot() <= slot {
		// Slot doesn't strictly belong to the epoch, create a clone with fixed values.
		epoch_index = epoch.clone_for_slot(slot).epoch_index;
	}

	let expected_author = secondary_slot_author(slot, &epoch.authorities, epoch.randomness)?;

	for (authority_id, authority_index) in keys {
		if authority_id == expected_author {
			let pre_digest = if author_secondary_vrf {
				let data = make_vrf_sign_data(&epoch.randomness, slot, epoch_index);
				let result =
					keystore.sr25519_vrf_sign(AuthorityId::ID, authority_id.as_ref(), &data);
				if let Ok(Some(vrf_signature)) = result {
					Some(PreDigest::SecondaryVRF(SecondaryVRFPreDigest {
						slot,
						authority_index: *authority_index as u32,
						vrf_signature,
					}))
				} else {
					None
				}
			} else if keystore.has_keys(&[(authority_id.to_raw_vec(), AuthorityId::ID)]) {
				Some(PreDigest::SecondaryPlain(SecondaryPlainPreDigest {
					slot,
					authority_index: *authority_index as u32,
				}))
			} else {
				None
			};

			if let Some(pre_digest) = pre_digest {
				return Some((pre_digest, authority_id.clone()))
			}
		}
	}

	None
}

/// Tries to claim the given slot number. This method starts by trying to claim
/// a primary VRF based slot. If we are not able to claim it, then if we have
/// secondary slots enabled for the given epoch, we will fallback to trying to
/// claim a secondary slot.
pub fn claim_slot(
	slot: Slot,
	epoch: &Epoch,
	keystore: &KeystorePtr,
) -> Option<(PreDigest, AuthorityId)> {
	let authorities = epoch
		.authorities
		.iter()
		.enumerate()
		.map(|(index, a)| (a.0.clone(), index))
		.collect::<Vec<_>>();
	claim_slot_using_keys(slot, epoch, keystore, &authorities)
}

/// Like `claim_slot`, but allows passing an explicit set of key pairs. Useful if we intend
/// to make repeated calls for different slots using the same key pairs.
pub fn claim_slot_using_keys(
	slot: Slot,
	epoch: &Epoch,
	keystore: &KeystorePtr,
	keys: &[(AuthorityId, usize)],
) -> Option<(PreDigest, AuthorityId)> {
	claim_primary_slot(slot, epoch, epoch.config.c, keystore, keys).or_else(|| {
		if epoch.config.allowed_slots.is_secondary_plain_slots_allowed() ||
			epoch.config.allowed_slots.is_secondary_vrf_slots_allowed()
		{
			claim_secondary_slot(
				slot,
				epoch,
				keys,
				keystore,
				epoch.config.allowed_slots.is_secondary_vrf_slots_allowed(),
			)
		} else {
			None
		}
	})
}

/// Claim a primary slot if it is our turn.  Returns `None` if it is not our turn.
/// This hashes the slot number, epoch, genesis hash, and chain randomness into
/// the VRF.  If the VRF produces a value less than `threshold`, it is our turn,
/// so it returns `Some(_)`. Otherwise, it returns `None`.
fn claim_primary_slot(
	slot: Slot,
	epoch: &Epoch,
	c: (u64, u64),
	keystore: &KeystorePtr,
	keys: &[(AuthorityId, usize)],
) -> Option<(PreDigest, AuthorityId)> {
	let mut epoch_index = epoch.epoch_index;
	if epoch.end_slot() <= slot {
		// Slot doesn't strictly belong to the epoch, create a clone with fixed values.
		epoch_index = epoch.clone_for_slot(slot).epoch_index;
	}

	let data = make_vrf_sign_data(&epoch.randomness, slot, epoch_index);

	for (authority_id, authority_index) in keys {
		let result = keystore.sr25519_vrf_sign(AuthorityId::ID, authority_id.as_ref(), &data);
		if let Ok(Some(vrf_signature)) = result {
			let threshold = calculate_primary_threshold(c, &epoch.authorities, *authority_index);

			let can_claim = authority_id
				.as_inner_ref()
				.make_bytes::<AUTHORING_SCORE_LENGTH>(
					AUTHORING_SCORE_VRF_CONTEXT,
					&data.as_ref(),
					&vrf_signature.pre_output,
				)
				.map(|bytes| u128::from_le_bytes(bytes) < threshold)
				.unwrap_or_default();

			if can_claim {
				let pre_digest = PreDigest::Primary(PrimaryPreDigest {
					slot,
					authority_index: *authority_index as u32,
					vrf_signature,
				});

				return Some((pre_digest, authority_id.clone()))
			}
		}
	}

	None
}

#[cfg(test)]
mod tests {
	use super::*;
	use sp_consensus_babe::{
		AllowedSlots, AuthorityId, BabeEpochConfiguration, Epoch, RANDOMNESS_LENGTH,
	};
	use sp_core::{crypto::Pair as _, sr25519::Pair};
	use sp_keystore::testing::MemoryKeystore;

	#[test]
	fn claim_secondary_plain_slot_works() {
		let keystore: KeystorePtr = MemoryKeystore::new().into();
		let valid_public_key = keystore
			.sr25519_generate_new(AuthorityId::ID, Some(sp_core::crypto::DEV_PHRASE))
			.unwrap();

		let authorities = vec![
			(AuthorityId::from(Pair::generate().0.public()), 5),
			(AuthorityId::from(Pair::generate().0.public()), 7),
		];

		let mut epoch = Epoch {
			epoch_index: 10,
			start_slot: 0.into(),
			duration: 20,
			authorities: authorities.clone(),
			randomness: Default::default(),
			config: BabeEpochConfiguration {
				c: (3, 10),
				allowed_slots: AllowedSlots::PrimaryAndSecondaryPlainSlots,
			},
		}
		.into();

		assert!(claim_slot(10.into(), &epoch, &keystore).is_none());

		epoch.authorities.push((valid_public_key.into(), 10));
		assert_eq!(claim_slot(10.into(), &epoch, &keystore).unwrap().1, valid_public_key.into());
	}

	#[test]
	fn secondary_slot_author_selection_works() {
		let authorities = (0..1000)
			.map(|i| (AuthorityId::from(Pair::generate().0.public()), i))
			.collect::<Vec<_>>();

		let randomness = [3; RANDOMNESS_LENGTH];

		assert_eq!(
			*secondary_slot_author(100.into(), &authorities, randomness).unwrap(),
			authorities[167].0
		);
	}
}