use curve25519_dalek::constants::RISTRETTO_BASEPOINT_POINT;
use curve25519_dalek::ristretto::RistrettoPoint;
use curve25519_dalek::scalar::Scalar;
use curve25519_dalek::traits::Identity;
use rand::rngs::ThreadRng;

pub type Secret = Scalar;

#[derive(Copy, Clone)]
pub struct Share {
    index: u32,
    value: Scalar,
}

pub struct Commitment(Vec<RistrettoPoint>);

/// Create secret shares for a given secret.
pub fn generate_shares(
    secret: Secret,
    numshares: u32,
    threshold: u32,
) -> Result<(Commitment, Vec<Share>), &'static str> {
    if threshold < 1 {
        return Err("Threshold cannot be 0");
    }
    if numshares < 1 {
        return Err("Number of shares cannot be 0");
    }
    if threshold > numshares {
        return Err("Threshold cannot exceed numshares");
    }

    let numcoeffs = (threshold - 1) as usize;

    let mut coefficients: Vec<Scalar> = Vec::with_capacity(numcoeffs);

    let mut rng: ThreadRng = rand::thread_rng();

    let mut shares: Vec<Share> = Vec::with_capacity(numshares as usize);

    let mut commitment = Commitment(Vec::with_capacity(threshold as usize));

    for _ in 0..numcoeffs {
        coefficients.push(Scalar::random(&mut rng));
    }

    for share_index in 1..numshares + 1 {
        // Evaluate the polynomial with `secret` as the constant term
        // and `coeffs` as the other coefficients at the point x=share_index
        // using Horner's method
        let scalar_index = Scalar::from(share_index);
        let mut value = Scalar::zero();
        for i in (0..numcoeffs).rev() {
            value += coefficients[i];
            value *= scalar_index;
        }
        value += secret;
        shares.push(Share {
            index: share_index,
            value: value,
        });
    }

    commitment.0.push(RISTRETTO_BASEPOINT_POINT * secret);
    for c in coefficients {
        commitment.0.push(RISTRETTO_BASEPOINT_POINT * c);
    }

    Ok((commitment, shares))
}

/// Verify that a share is consistent with a commitment.
pub fn verify_share(share: &Share, commitment: &Commitment) -> Result<(), &'static str> {
    let f_result = RISTRETTO_BASEPOINT_POINT * share.value;

    let x = Scalar::from(share.index);

    let (_, result) = commitment.0.iter().fold(
        (Scalar::one(), RistrettoPoint::identity()),
        |(x_to_the_i, sum_so_far), comm_i| (x_to_the_i * x, sum_so_far + x_to_the_i * comm_i),
    );

    if f_result == result {
        Ok(())
    } else {
        Err("Share is inconsistent")
    }
}

/// Reconstruct the secret from enough (at least the threshold) already-verified shares.
pub fn reconstruct_secret(shares: &Vec<Share>) -> Result<Secret, &'static str> {
    let numshares = shares.len();

    if numshares < 1 {
        return Err("No shares provided");
    }

    let mut lagrange_coeffs: Vec<Scalar> = Vec::with_capacity(numshares);

    for i in 0..numshares {
        let mut num = Scalar::one();
        let mut den = Scalar::one();
        for j in 0..numshares {
            if j == i {
                continue;
            }
            num *= Scalar::from(shares[j].index);
            den *= Scalar::from(shares[j].index) - Scalar::from(shares[i].index);
        }
        if den == Scalar::zero() {
            return Err("Duplicate shares provided");
        }
        lagrange_coeffs.push(num * den.invert());
    }

    let mut secret = Scalar::zero();

    for i in 0..numshares {
        secret += lagrange_coeffs[i] * shares[i].value;
    }

    Ok(secret)
}

/// Create a proactive update.
pub fn create_update(
    num_shares: u32,
    threshold: u32,
) -> Result<(Commitment, Vec<Share>), &'static str> {
    generate_shares(Scalar::zero(), num_shares, threshold)
}

/// Apply the commitment for the update to the master commitment.
pub fn apply_commitment_update(
    old_commitment: &Commitment,
    update: &Commitment,
) -> Result<Commitment, &'static str> {
    let mut new_commitments = Commitment(Vec::with_capacity(old_commitment.0.len()));
    for i in 0..old_commitment.0.len() {
        let new_commitment = old_commitment.0[i] + update.0[i];
        new_commitments.0.push(new_commitment);
    }

    Ok(new_commitments)
}

/// Apply the share update to an existing share
pub fn apply_share_update(
    old_share: &Share,
    update: &Share,
    updated_commitment: &Commitment,
) -> Result<Share, &'static str> {
    let updated_share = Share {
        index: old_share.index,
        value: old_share.value + update.value,
    };

    match verify_share(&updated_share, updated_commitment) {
        Ok(n) => n,
        Err(_) => return Err("Error when validating share"),
    };

    Ok(updated_share)
}

#[cfg(test)]
mod tests {
    use crate::vss::*;

    #[test]
    fn share_simple() {
        let s = Secret::from(42u32);

        let res = generate_shares(s, 5, 2);
        assert!(res.is_ok());
        let (com, shares) = res.unwrap();
        assert!(shares.len() == 5);
        assert!(com.0.len() == 2);

        let mut recshares: Vec<Share> = Vec::new();
        recshares.push(shares[1]);
        recshares.push(shares[3]);
        let recres = reconstruct_secret(&recshares);
        assert!(recres.is_ok());
        assert_eq!(recres.unwrap(), s);
    }

    #[test]
    fn share_not_enough() {
        let s = Secret::from(42u32);

        let res = generate_shares(s, 5, 2);
        assert!(res.is_ok());
        let (_, shares) = res.unwrap();

        let mut recshares: Vec<Share> = Vec::new();
        recshares.push(shares[1]);
        let recres = reconstruct_secret(&recshares);
        assert!(recres.is_ok());
        assert_ne!(recres.unwrap(), s);
    }

    #[test]
    fn share_dup() {
        let s = Secret::from(42u32);

        let res = generate_shares(s, 5, 2);
        assert!(res.is_ok());
        let (_, shares) = res.unwrap();

        let mut recshares: Vec<Share> = Vec::new();
        recshares.push(shares[1]);
        recshares.push(shares[1]);
        let recres = reconstruct_secret(&recshares);
        assert!(recres.is_err());
        assert_eq!(recres.err(), Some("Duplicate shares provided"));
    }

    #[test]
    fn share_badparams() {
        let s = Secret::from(42u32);

        {
            let res = generate_shares(s, 5, 0);
            assert!(res.is_err());
            assert_eq!(res.err(), Some("Threshold cannot be 0"));
        }
        {
            let res = generate_shares(s, 0, 3);
            assert!(res.is_err());
            assert_eq!(res.err(), Some("Number of shares cannot be 0"));
        }
        {
            let res = generate_shares(s, 1, 3);
            assert!(res.is_err());
            assert_eq!(res.err(), Some("Threshold cannot exceed numshares"));
        }
    }

    #[test]
    fn share_commitment_valid() {
        let s = Secret::from(42u32);

        let res = generate_shares(s, 8, 3);
        assert!(res.is_ok());
        let (com, shares) = res.unwrap();

        for share in shares {
            let is_valid = verify_share(&share, &com);
            assert!(is_valid.is_ok());
        }
    }

    #[test]
    fn share_commitment_invalid() {
        let s1 = Secret::from(42u32);
        let s2 = Secret::from(42u32);

        let res1 = generate_shares(s1, 8, 3);
        assert!(res1.is_ok());
        let (_, shares1) = res1.unwrap();

        let res2 = generate_shares(s2, 8, 3);
        assert!(res2.is_ok());
        let (com2, _) = res2.unwrap();

        for share in shares1 {
            // test that commitments to a different set of shares fails
            let is_valid = verify_share(&share, &com2);
            assert!(is_valid.is_err());
        }
    }

    #[test]
    fn share_update_valid() {
        let s = Secret::from(42u32);

        let res = generate_shares(s, 5, 2);
        assert!(res.is_ok());
        let (com, shares) = res.unwrap();

        // Create a new update
        let update_comm_res = create_update(5, 2);
        assert!(update_comm_res.is_ok());
        let (com_update, shares_update) = update_comm_res.unwrap();

        // Apply the update to previously-created committment
        let updated_commitment_res = apply_commitment_update(&com, &com_update);
        assert!(updated_commitment_res.is_ok());
        let updated_commitment = updated_commitment_res.unwrap();

        // Distribute the update to all owners of shares
        // Verify and apply the update
        // Collect updated shares to test secret reconstruction after
        let mut updates_to_shares = shares_update.iter();
        let mut updated_shares: Vec<Share> = Vec::with_capacity(5);
        for share in &shares {
            let share_update = updates_to_shares.next().unwrap();
            let update_share_res = apply_share_update(&share, &share_update, &updated_commitment);
            assert!(update_share_res.is_ok());
            let updated_share = update_share_res.unwrap();
            updated_shares.push(updated_share);
        }

        // assert that we can recover the original secret using the updated shares
        let recres = reconstruct_secret(&updated_shares);
        assert!(recres.is_ok());
        assert_eq!(recres.unwrap(), s);

        // test that the old shares are not valid with the updated commitment
        for share in &shares {
            let is_valid = verify_share(&share, &updated_commitment);
            assert!(is_valid.is_err());
        }
    }

    #[test]
    fn share_update_valid_multiple_updates() {
        let s = Secret::from(42u32);

        let res = generate_shares(s, 5, 2);
        assert!(res.is_ok());
        let (com, shares) = res.unwrap();

        // final shares to test at the end
        let mut updated_shares: Vec<Share> = Vec::with_capacity(5);

        // override these each time in the loop
        let mut next_shares = shares;
        let mut next_com = com;

        for i in 0..5 {
            // Create a new update
            let update_comm_res = create_update(5, 2);
            assert!(update_comm_res.is_ok());
            let (com_update, shares_update) = update_comm_res.unwrap();

            // Apply the update to previously-created committment
            let updated_commitment_res = apply_commitment_update(&next_com, &com_update);
            assert!(updated_commitment_res.is_ok());
            let updated_commitment = updated_commitment_res.unwrap();

            let mut iterim_shares: Vec<Share> = Vec::with_capacity(5);
            let mut updates_to_shares = shares_update.iter();
            for share in &next_shares {
                let share_update = updates_to_shares.next().unwrap();
                let update_share_res =
                    apply_share_update(&share, &share_update, &updated_commitment);

                assert!(update_share_res.is_ok());
                let updated_share = update_share_res.unwrap();

                // only collect the last round of updated shares
                if i == 4 {
                    updated_shares.push(updated_share);
                }
                iterim_shares.push(updated_share);
            }
            next_shares = iterim_shares;
            next_com = updated_commitment;
        }

        // assert that we can recover the original secret using the updated shares
        let recres = reconstruct_secret(&updated_shares);
        assert!(recres.is_ok());
        assert_eq!(recres.unwrap(), s);
    }
}