// SPDX-License-Identifier: Apache-2.0

//! This module implements cryptographic primitives for verifiable secret sharing (VSS).
//! Each share contains a (Merkle) commitment to the set of all shares generated by the
//! sharing algorithm, so that each share can be individually verified during reconstruction.

use rand_chacha::rand_core::SeedableRng;

use super::channel::{decrypt_message, encrypt_message};

mod shamir;
mod utils;

#[allow(non_upper_case_globals)]
const λ_bits: usize = 256;

#[allow(non_upper_case_globals)]
const λ: usize = λ_bits / 8;

const MERKLE_TREE_DEPTH: u32 = 7;

/// Encodes a VSS share.
#[derive(Clone)]
pub struct VSSShare {
    /// we use the x-coordinate to uniquely identify shares
    pub x: Vec<u8>,
    /// we use the y-coordinate as the share
    pub y: Vec<u8>,
    /// AES encryption of the secret message
    pub encrypted_secret: Vec<u8>,
    /// Merkle-root commitment to all shares
    pub commitment: Vec<u8>,
    /// bottom-up Merkle authentication path
    /// bool denotes isLeft, while vec<u8> is the SHA-384 hash
    pub merkle_path: Vec<(bool, Vec<u8>)>,
}

/// Custom error type for Verifiable Secret Sharing (VSS) operations.
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum DerecVSSError {
    #[error("inconsistent ciphertexts")]
    InconsistentCiphertexts,

    #[error("inconsistent commitments")]
    InconsistentCommitments,

    #[error("one or more shares are corrupted")]
    CorruptShares,

    #[error("insufficient shares")]
    InsufficientShares,

    #[error("decryption failed")]
    DecryptionFailure,

    #[error("invalid access structure")]
    InvalidAccessStructure,
}

/// Generates VSS shares for a given secret using Shamir's Secret Sharing scheme,
/// with additional cryptographic commitments for verifiability during reconstruction.
///
/// # Arguments
///
/// * `access_structure` - A tuple `(t, n)` where `t` is the threshold number of shares required to reconstruct the secret,
///   and `n` is the total number of shares to generate. Must satisfy `2 <= t <= n` and `n <= 128`.
/// * `msg` - The secret message to be shared, as a byte slice.
/// * `rand` - A cryptographically secure random seed of length `λ` (32 bytes).
///
/// # Returns
///
/// A vector of `VSSShare` structs, each containing:
///   - The share's x and y coordinates (as bytes),
///   - An AES-encrypted ciphertext of the secret data,
///   - A Merkle root commitment to all shares,
///   - A Merkle authentication path for the share.
///
/// # Details
///
/// - The function derives a pseudo-random AES key and nonce from the message and random seed.
/// - The message is encrypted using AES with the derived key and nonce.
/// - The AES key is split into shares using Shamir's Secret Sharing.
/// - A Merkle tree is constructed over the shares for verifiable commitments.
/// - Each share includes its Merkle authentication path for individual verification.
///
/// # Example
///
/// ```rust
/// use derec_cryptography::vss::{share, recover, VSSShare};
/// let shares = share((3, 5), b"my secret", &[0u8; 32]).unwrap();
/// assert_eq!(shares.len(), 5);
/// ```
pub fn share(
    access_structure: (u64, u64), // (t, n)
    msg: &[u8],
    entropy: &[u8; λ],
) -> Result<Vec<VSSShare>, DerecVSSError> {
    if (access_structure.0 > access_structure.1) || (access_structure.0 < 2) {
        return Err(DerecVSSError::InvalidAccessStructure);
    }

    // we can only support up to 2^7 = 128 shares
    if access_structure.1 > 1 << MERKLE_TREE_DEPTH {
        return Err(DerecVSSError::InvalidAccessStructure);
    }

    //pseudo-random key derivation
    let hash = utils::random_oracle(msg, entropy, &[]);
    let k: [u8; λ] = hash[..λ].try_into().unwrap();
    let nonce: [u8; λ] = hash[λ..2 * λ].try_into().unwrap();
    let seed1: [u8; λ] = hash[2 * λ..3 * λ].try_into().unwrap();
    let seed2: [u8; λ] = hash[3 * λ..4 * λ].try_into().unwrap();

    //AES encrypt the message using the pseudo-random key k
    let c = encrypt_message(msg, &k, &nonce).unwrap();

    // generate shares of the AES key k
    let shamir_shares = shamir::share(
        &k,
        access_structure,
        &mut rand_chacha::ChaCha8Rng::from_seed(seed1),
    );

    let merkle_tree = utils::build_merkle_tree(
        &shamir_shares,
        MERKLE_TREE_DEPTH,
        &mut rand_chacha::ChaCha8Rng::from_seed(seed2),
    );
    let merkle_proofs =
        utils::extract_merkle_proofs(&merkle_tree, MERKLE_TREE_DEPTH, access_structure.1);

    let mut output = vec![];
    for (i, (x, y)) in shamir_shares.iter().enumerate() {
        output.push(VSSShare {
            x: x.to_owned(),
            y: y.to_owned(),
            encrypted_secret: c.clone(),
            commitment: merkle_tree[0].clone(),
            merkle_path: merkle_proofs[i].to_owned(),
        });
    }
    Ok(output)
}

/// Recovers the secret-shared data from a set of VSS shares.
///
/// This function attempts to reconstruct the secret by first verifying the integrity and consistency
/// of the provided shares using Merkle proofs, with respect to a common Merkle commitment.
/// If no errors are detected, it reconstructs the AES key using Shamir interpolation and then
/// proceeds to decrypt the ciphertext using the recovered AES key.
///
/// # Arguments
///
/// * `shares` - A vector of `VSSShare` structs, each containing the necessary data for reconstruction,
///   including the share's coordinates, encrypted secret, Merkle commitment, and authentication path.
///
/// # Returns
///
/// * `Ok(Vec<u8>)` containing the recovered secret message if reconstruction and decryption succeed.
/// * `Err(VSSError)` if shares are inconsistent, corrupted, insufficient, or decryption fails.
///
/// # Errors
///
/// Returns a `VSSError` variant in the following cases:
/// - `VSSError::CorruptShares` if one or more shares are detected as corrupted.
/// - `VSSError::InconsistentCommitments` if Merkle commitments do not match.
/// - `VSSError::InsufficientShares` if not enough valid shares are provided for reconstruction.
/// - `VSSError::DecryptionFailure` if the reconstructed key fails to decrypt the secret.
///
/// # Example
///
/// ```rust
/// use derec_cryptography::vss::{share, recover, VSSShare};
/// let shares = share((3, 5), b"my secret", &[0u8; 32]).unwrap();
/// let secret = recover(&shares[..3].to_vec()).unwrap();
/// assert_eq!(secret, b"my secret");
/// ```
pub fn recover(shares: &Vec<VSSShare>) -> Result<Vec<u8>, DerecVSSError> {
    assert!(!shares.is_empty());

    let detected_error = utils::detect_error(shares);
    if let Some(err) = detected_error {
        return Err(err);
    }

    // no error detected so far, let's try shamir reconstruction
    let shamir_shares = shares.iter().map(|s| (s.x.clone(), s.y.clone())).collect();
    let k = shamir::recover(shamir_shares);

    // let's attempt to decrypt using the shamir-reconstructed key
    let c = shares[0].encrypted_secret.clone();
    decrypt_message(&c, &k).map_err(|_| DerecVSSError::InsufficientShares)
}
