//
// Copyright 2020 Signal Messenger, LLC.
// SPDX-License-Identifier: AGPL-3.0-only
//
use std::borrow::Cow;
use std::collections::HashMap;

// POKSHO implements the "Sigma protocol for arbitrary linear relations" described in section
// 19.5.3 of https://crypto.stanford.edu/~dabo/cryptobook/BonehShoup_0_4.pdf
//
// We adopt the view that we are proving knowledge of the preimage of a group homomorphism
// from groups G1 -> G2, where elements in G1 are vectors of scalars and elements in G2 are vectors of
// Ristretto "points".  The homomomorphism can be viewed as a system of equations of the form:
//
// P = sP + sP + sP + ...
// P = sP + sP + sP + ...
// P = sP + sP + sP + ...
// ...
//
// Where s are any scalars and P are any points.  The left-hand side of these equations describes
// the element in G2 which is the image of the homomorphism, and the scalars form an element in G1
// which we are proving knowledge of (this preimage in G1 can be viewed as a vector if the system
// of equations is rewritten as matrix multiplication).
//
// We use the term "statement" for the above system of equations written with both points and
// scalars as variables.  For example, the statement for signatures is the single equation:
//
// A = a*G
//
// where "A" is the public key, "a" is the private key, and "G" is the base point.  The
// discrete-log equality statement used in VRFs could be written as:
//
// A = a*G
// B = a*H
//
// We use the term "witness" for a vector of scalars that satisfies a given statement and a given
// set of point assignments (equivalently: a witness is an element in G1 that is a preimage of the
// left-hand-side under the group homomorphism).  In the above cases, "a" is the witness.
//
// The zero-knowledge proof is a standard Fiat-Shamir Sigma/Schnorr proof of knowledge.  To
// implement Fiat-Shamir hashing we use the SHO/HMAC-SHA256 construct, which provides a stateful object
// which we can use to Absorb data, and then Squeeze out arbitrary-length output.  We use a SHO
// object not only to produce the Schnorr challenge, but also to produce the Schnorr nonce by
// hashing some caller-supplied random data, the witness, and the message.  This "synthetic nonce"
// strategy is intended to ensure that the nonce appears random to any attacker and that different
// Schnorr challenges will never be used with the same nonce.
//
// Below we describe the hash inputs to SHO/HMAC-SHA256:
//
//  L : bytes, label = "POKSHO_Ristretto_SHOHMACSHA256"
//  D : bytes, description of statement - see below
//  a : G1, witness scalars for statement
//  A : G2, point values for statement = homomorphism(a)
//  Z : bytes, random = 32 byes of randomness
//  M : bytes, message to be signed, if any
//  r : G1, Schnorr nonce
//  R : G2, Schnorr commitment = homomorphism(r)
//  h : integer, Schnorr challenge

//  sho = SHO(L)
//  sho.AbsorbAndRatchet(D || A)
//         sho2 = sho.Clone()
//         sho2.AbsorbAndRatchet(Z || a)
//         sho2.AbsorbAndRatchet(M)
//         r = sho2.Squeeze(64 * num_scalars)
//  sho.AbsorbAndRatchet(R || M)
//  h = Squeeze(64)
//
//  Description format (D)
//  ---
//  Ne : number of equations (1-255)
//  for i=1..Ne:
//    point_index : 0..255 (0 = base point)
//    Nt : number of terms (1-255)
//    for j=1..Nt:
//      scalar_index: 0-255
//      point_index: 0-255 (0 = base point)
//
//  Point values for statement (A) and commitment (R)
//  ---
//   for index=1..total number of points (excluding base point at index 0):
//     RistrettoPoint
//
//  Witness (a)
//  ---
//  for index=0..total number of scalars:
//   RistrettoScalar
use PokshoError::*;
use curve25519_dalek::constants::RISTRETTO_BASEPOINT_POINT;
use curve25519_dalek::ristretto::RistrettoPoint;
use curve25519_dalek::scalar::Scalar;
use curve25519_dalek::traits::MultiscalarMul;

use crate::args::*;
use crate::errors::*;
use crate::proof::*;
use crate::shoapi::{ShoApi, ShoApiExt as _};
use crate::shohmacsha256::ShoHmacSha256;
use crate::simple_types::*;

type ScalarIndex = u8;
type PointIndex = u8;

struct Term {
    scalar: ScalarIndex,
    point: PointIndex,
}

struct Equation {
    lhs: PointIndex,
    rhs: Vec<Term>,
}

pub struct Statement {
    // We store the Schnorr ZKP equations using scalar and point indices
    // which are numbered from zero, and are assigned sequentially based
    // on the order in which the point or scalar appears in the equations
    // (except point index 0 is pre-assigned to "G", the Ristretto base point)
    //
    // We also store maps from string names -> indices, and vectors for
    // the reverse map.  The former map is used when adding new equations,
    // and the latter is used when instantiating these indices with
    // concrete values.
    equations: Vec<Equation>,
    scalar_map: HashMap<Cow<'static, str>, ScalarIndex>,
    scalar_vec: Vec<Cow<'static, str>>,
    point_map: HashMap<Cow<'static, str>, PointIndex>,
    point_vec: Vec<Cow<'static, str>>,
}

impl Statement {
    pub fn new() -> Self {
        let mut point_map = HashMap::new();
        point_map.insert("G".into(), 0); // G is base point
        let point_vec = vec!["G".into()];
        Statement {
            equations: Vec::new(),
            scalar_map: HashMap::new(),
            scalar_vec: Vec::new(),
            point_map,
            point_vec,
        }
    }

    // panics on invalid input
    pub fn add(&mut self, lhs_str: &str, rhs_pairs: &[(&str, &str)]) {
        if (lhs_str.is_empty())
            || (rhs_pairs.is_empty())
            || (rhs_pairs.len() > 255)
            || (self.equations.len() >= 255)
        {
            panic!("Unexpected input sizes to add");
        }
        let lhs = self
            .add_point(lhs_str.to_string())
            .expect("add_point succeeds");
        let mut rhs = Vec::<Term>::with_capacity(rhs_pairs.len());
        for pair in rhs_pairs {
            if pair.0.is_empty() || pair.1.is_empty() {
                panic!("Unexpected pair size");
            }
            let scalar = self
                .add_scalar(pair.0.to_string())
                .expect("add_scalar succeeds");
            let point = self
                .add_point(pair.1.to_string())
                .expect("add_point succeeds");
            rhs.push(Term { scalar, point });
        }
        self.equations.push(Equation { lhs, rhs });
    }

    pub fn prove(
        &self,
        scalar_args: &ScalarArgs,
        point_args: &PointArgs,
        message: &[u8],
        randomness: &[u8], // must be 32 bytes
    ) -> Result<Vec<u8>, PokshoError> {
        if randomness.len() != 32 {
            return Err(PokshoError::BadArgs);
        }
        let g1 = self.sort_scalars(scalar_args)?;
        let all_points = self.sort_points(point_args)?;

        // Absorb the protocol label L, description of statement D, and point values A
        let mut sho = ShoHmacSha256::new(b"POKSHO_Ristretto_SHOHMACSHA256"); // L
        sho.absorb(&self.to_bytes()); // D
        for point in &all_points {
            // A
            sho.absorb(&point.compress().to_bytes());
        }

        sho.ratchet(); // Ratchet

        // Random nonce
        // "Synthetic" nonce based on hashing randomness, witness (private scalars) and message
        let mut sho2 = sho.clone();
        sho2.absorb(randomness); // Z
        for scalar in &g1 {
            sho2.absorb(&scalar.to_bytes()); // a
        }
        sho2.ratchet(); // Ratchet
        sho2.absorb_and_ratchet(message); // M
        let blinding_scalar_bytes = sho2.squeeze_and_ratchet(g1.len() * 64);

        let nonce: G1 = blinding_scalar_bytes
            .as_chunks::<64>()
            .0
            .iter()
            .map(Scalar::from_bytes_mod_order_wide)
            .collect();

        // Commitment from nonce by applying homomorphism F: commitment = F(nonce)
        let commitment = self.homomorphism_with_subtraction(&nonce, &all_points, None);

        // Challenge from commitment and message
        for point in &commitment {
            sho.absorb(&point.compress().to_bytes());
        }
        sho.absorb_and_ratchet(message);
        let challenge = Scalar::from_bytes_mod_order_wide(&sho.squeeze_and_ratchet_as_array());

        // Response
        let response = nonce
            .into_iter()
            .zip(g1)
            .map(|(nonce, g1)| nonce + (g1 * challenge))
            .collect();

        let proof = Proof {
            challenge,
            response,
        };

        // Verify before returning, since a bad proof could indicate
        // a glitched/faulty response that leaks private keys, or incorrect inputs
        let proof_bytes = proof.to_bytes();
        match self.verify_proof(&proof_bytes, point_args, message) {
            Err(VerificationFailure) => Err(ProofCreationVerificationFailure),
            Err(e) => Err(e),
            Ok(_) => Ok(proof_bytes),
        }
    }

    pub fn verify_proof(
        &self,
        proof_bytes: &[u8],
        point_args: &PointArgs,
        message: &[u8],
    ) -> Result<(), PokshoError> {
        let proof = Proof::from_slice(proof_bytes).ok_or(VerificationFailure)?;
        if proof.response.len() != self.scalar_vec.len() {
            return Err(VerificationFailure);
        }
        let all_points = self.sort_points(point_args)?;

        // Absorb the protocol label L, statement description D, and point values A
        let mut sho = ShoHmacSha256::new(b"POKSHO_Ristretto_SHOHMACSHA256"); // L
        sho.absorb(&self.to_bytes()); // D
        for point in &all_points {
            // A
            sho.absorb(&point.compress().to_bytes());
        }
        sho.ratchet();

        // Reconstruct commitment
        //
        // commitment R = F(s) - h*A
        //
        // F: homomorphism
        // s: response element in G1
        // h: challenge scalar
        // A: element in G2 whose preimage we are proving knowledge of (i.e. LHS of Schnorr eqns)
        let commitment =
            self.homomorphism_with_subtraction(&proof.response, &all_points, Some(proof.challenge));

        // Reconstruct challenge from commitment and message
        for point in &commitment {
            // R
            sho.absorb(&point.compress().to_bytes());
        }
        sho.absorb_and_ratchet(message); // M
        let challenge = Scalar::from_bytes_mod_order_wide(&sho.squeeze_and_ratchet_as_array());

        // Check challenge (const time)
        if challenge == proof.challenge {
            Ok(())
        } else {
            Err(VerificationFailure)
        }
    }

    fn add_scalar(
        &mut self,
        scalar_name: impl Into<Cow<'static, str>>,
    ) -> Result<ScalarIndex, PokshoError> {
        let scalar_name = scalar_name.into();
        match self.scalar_map.get(&scalar_name) {
            Some(index) => Ok(*index),
            None => {
                assert!(self.scalar_map.len() == self.scalar_vec.len());
                let Ok(new_index) = self.scalar_map.len().try_into() else {
                    return Err(BadArgs);
                };
                self.scalar_map.insert(scalar_name.clone(), new_index);
                self.scalar_vec.push(scalar_name.clone());
                Ok(new_index)
            }
        }
    }

    fn add_point(
        &mut self,
        point_name: impl Into<Cow<'static, str>>,
    ) -> Result<PointIndex, PokshoError> {
        let point_name = point_name.into();
        match self.point_map.get(&point_name) {
            Some(index) => Ok(*index),
            None => {
                assert!(self.point_map.len() == self.point_vec.len());
                let Ok(new_index) = self.point_map.len().try_into() else {
                    return Err(BadArgs);
                };
                self.point_map.insert(point_name.clone(), new_index);
                self.point_vec.push(point_name.clone());
                Ok(new_index)
            }
        }
    }

    fn to_bytes(&self) -> Vec<u8> {
        let equation_count =
            u8::try_from(self.equations.len()).expect("number of equations fits in a byte");
        let scalar_count =
            u8::try_from(self.scalar_map.len()).expect("number of scalars fits in a byte");
        let point_count =
            u8::try_from(self.point_map.len()).expect("number of points fits in a byte");
        let mut v = vec![equation_count];
        for Equation { lhs, rhs } in &self.equations {
            assert!(*lhs <= point_count);
            v.push(*lhs);
            let term_count = u8::try_from(rhs.len()).expect("number of terms fits in a byte");
            v.push(term_count);
            for Term { scalar, point } in rhs {
                assert!(*scalar < scalar_count);
                assert!(*point < point_count);
                v.push(*scalar);
                v.push(*point);
            }
        }
        v
    }

    // Applies the homomorphism from G1 -> G2
    // If given a challenge h, also subtracts h*A for efficient recovery of
    // the Schnorr commitment
    fn homomorphism_with_subtraction(
        &self,
        g1: &[Scalar],
        all_points: &[RistrettoPoint],
        challenge: Option<Scalar>,
    ) -> G2 {
        self.equations
            .iter()
            .map(|e| {
                let scalar_iter = e
                    .rhs
                    .iter()
                    .map(|Term { scalar, point: _ }| g1[*scalar as usize]);
                let point_iter = e
                    .rhs
                    .iter()
                    .map(|Term { scalar: _, point }| all_points[*point as usize]);

                let (v_scalar, v_point) =
                    challenge.map(|h| (-h, all_points[e.lhs as usize])).unzip();

                let scalar_iter = scalar_iter.chain(v_scalar);
                let point_iter = point_iter.chain(v_point);

                // Could use vartime_multiscalar_mul in some cases, but in the
                // general case points might be secret (not just scalars!)
                RistrettoPoint::multiscalar_mul(scalar_iter, point_iter)
            })
            .collect()
    }

    fn sort_scalars(&self, scalar_args: &ScalarArgs) -> Result<G1, PokshoError> {
        if scalar_args.0.len() != self.scalar_vec.len() {
            return Err(BadArgsWrongNumberOfScalarArgs);
        }

        self.scalar_vec
            .iter()
            .map(|scalar_name| {
                scalar_args
                    .0
                    .get(scalar_name)
                    .copied()
                    .ok_or(BadArgsMissingScalarArg)
            })
            .collect()
    }

    fn sort_points(&self, point_args: &PointArgs) -> Result<Vec<RistrettoPoint>, PokshoError> {
        if point_args.0.len() != self.point_vec.len() - 1 {
            return Err(BadArgsWrongNumberOfPointArgs);
        }
        let try_iter_points = self.point_vec[1..].iter().map(|point_name| {
            point_args
                .0
                .get(point_name)
                .copied()
                .ok_or(BadArgsMissingPointArg)
        });

        [Ok(RISTRETTO_BASEPOINT_POINT)]
            .into_iter()
            .chain(try_iter_points)
            .collect()
    }
}

impl Default for Statement {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    #![allow(non_snake_case)]

    use super::*;

    #[test]
    fn test_statement_encoding() {
        let mut s = Statement::new();
        s.add("A", &[("a", "G")]);
        assert!(s.to_bytes() == vec![1, 1, 1, 0, 0]);

        let mut s = Statement::new();
        s.add("A", &[("a", "G")]);
        s.add("B", &[("a", "H")]);
        assert!(s.to_bytes() == vec![2, 1, 1, 0, 0, 2, 1, 0, 3]);

        let mut s = Statement::new();
        s.add("A", &[("a", "G"), ("b", "H")]);
        assert!(s.to_bytes() == vec![1, 1, 2, 0, 0, 1, 2]);
    }

    #[test]
    #[allow(
        clippy::cast_possible_truncation,
        clippy::needless_range_loop,
        clippy::redundant_clone,
        clippy::unwrap_used
    )]
    fn test_complex_statement() {
        let mut block32 = [0u8; 32];
        let mut block64a = [0u8; 64];
        let mut block64b = [0u8; 64];
        let mut block64c = [0u8; 64];
        let mut block64d = [0u8; 64];
        let mut block64h = [0u8; 64];
        let mut block64i = [0u8; 64];
        let block0 = [0u8; 0];
        for i in 0..32 {
            block32[i] = i as u8;
        }
        for i in 0..64 {
            block64a[i] = 10 + i as u8;
        }
        for i in 0..64 {
            block64b[i] = 20 + i as u8;
        }
        for i in 0..64 {
            block64c[i] = 30 + i as u8;
        }
        for i in 0..64 {
            block64d[i] = 40 + i as u8;
        }
        for i in 0..64 {
            block64h[i] = 50 + i as u8;
        }
        for i in 0..64 {
            block64i[i] = 60 + i as u8;
        }

        let randomness = block32;
        let message = block0;

        let scalar_bytes_a = block64a;
        let scalar_bytes_b = block64b;
        let scalar_bytes_c = block64c;
        let scalar_bytes_d = block64d;

        let scalar_bytes_h_point = block64h;
        let scalar_bytes_i_point = block64i;

        let a = Scalar::from_bytes_mod_order_wide(&scalar_bytes_a);
        let b = Scalar::from_bytes_mod_order_wide(&scalar_bytes_b);
        let c = Scalar::from_bytes_mod_order_wide(&scalar_bytes_c);
        let d = Scalar::from_bytes_mod_order_wide(&scalar_bytes_d);
        let H =
            Scalar::from_bytes_mod_order_wide(&scalar_bytes_h_point) * RISTRETTO_BASEPOINT_POINT;
        let I =
            Scalar::from_bytes_mod_order_wide(&scalar_bytes_i_point) * RISTRETTO_BASEPOINT_POINT;

        let A = a * RISTRETTO_BASEPOINT_POINT + b * H + c * I;
        let B = c * H + d * I;

        let mut st = Statement::new();
        st.add("A", &[("a", "G"), ("b", "H"), ("c", "I")]);
        st.add("B", &[("c", "H"), ("d", "I")]);
        assert!(st.to_bytes() == vec![2, 1, 3, 0, 0, 1, 2, 2, 3, 4, 2, 2, 2, 3, 3]);

        let mut scalar_args = ScalarArgs::new();
        scalar_args.add("a", a);
        scalar_args.add("b", b);
        scalar_args.add("c", c);
        scalar_args.add("d", d);
        let mut point_args = PointArgs::new();
        point_args.add("A", A);
        point_args.add("B", B);
        point_args.add("H", H);
        point_args.add("I", I);

        let mut scalar_args2 = scalar_args.clone();
        scalar_args2.add("abc", a);

        // Test bad args - extra scalar
        assert!(matches!(
            st.prove(&scalar_args2, &point_args, &message, &randomness),
            Err(PokshoError::BadArgsWrongNumberOfScalarArgs)
        ));

        // Good proof
        let mut proof = st
            .prove(&scalar_args, &point_args, &message, &randomness)
            .unwrap();
        st.verify_proof(&proof, &point_args, &message).unwrap();
        /*
        for b in proof.iter() {
            print!("0x{:02x}, ", b);
        }
        println!("");
        */
        assert!(
            proof
                == vec![
                    0x8e, 0xfc, 0x67, 0x6c, 0x33, 0xe6, 0xb2, 0xd0, 0x67, 0x0e, 0xd5, 0x46, 0x1a,
                    0x50, 0x7f, 0x6a, 0x4b, 0xc9, 0x15, 0x3e, 0x26, 0x1d, 0xb8, 0x0f, 0xa4, 0x38,
                    0xf3, 0xcd, 0x80, 0xa5, 0xc9, 0x09, 0xb1, 0x13, 0xcc, 0x0d, 0x79, 0x90, 0xad,
                    0x61, 0x6d, 0x0a, 0x2f, 0xc4, 0xb8, 0x31, 0xd0, 0x63, 0x57, 0xa5, 0xee, 0x5d,
                    0x36, 0xd4, 0x4b, 0x34, 0x27, 0xc7, 0x90, 0x10, 0x61, 0x18, 0x0c, 0x0f, 0xb1,
                    0x79, 0x8c, 0x51, 0x68, 0x0f, 0xe2, 0x1b, 0x9f, 0x98, 0xe9, 0x79, 0x55, 0xb1,
                    0x59, 0x7c, 0x49, 0x31, 0x47, 0x25, 0xc1, 0x54, 0x6a, 0x36, 0x93, 0x28, 0xcf,
                    0x54, 0xda, 0xae, 0x71, 0x0b, 0xfc, 0x4a, 0x99, 0x11, 0x42, 0x2a, 0xa7, 0x7e,
                    0xd6, 0xd7, 0x23, 0x1d, 0xe3, 0x00, 0x3b, 0xa5, 0xae, 0x9d, 0x9f, 0xd0, 0xc5,
                    0x3c, 0xed, 0x7a, 0xd7, 0x82, 0xe2, 0x9b, 0x04, 0x68, 0x4a, 0x07, 0x22, 0x1a,
                    0x6e, 0xf4, 0x7c, 0xe6, 0x1d, 0x81, 0x7f, 0x01, 0x11, 0x7c, 0xf5, 0x9d, 0xf6,
                    0x9a, 0xc3, 0x5b, 0x5b, 0xb5, 0x90, 0xf1, 0xf7, 0xb6, 0xd0, 0x29, 0x71, 0x7b,
                    0xc1, 0xa6, 0x25, 0x01,
                ]
        );

        // Test bad args - extra point
        let mut point_args2 = point_args.clone();
        point_args2.add("xyz", A);
        assert!(matches!(
            st.verify_proof(&proof, &point_args2, &message),
            Err(PokshoError::BadArgsWrongNumberOfPointArgs)
        ));

        // Test bad message
        assert!(matches!(
            st.verify_proof(&proof, &point_args, &block32),
            Err(VerificationFailure)
        ));

        // Test bad proof #1 - extra byte at end
        let mut proof2 = proof.clone();
        proof2.push(0);
        assert!(matches!(
            st.verify_proof(&proof2, &point_args, &message),
            Err(VerificationFailure)
        ));

        // Test bad proof #2 - last byte changed
        let prooflen = proof.len();
        proof[prooflen - 1] += 1;
        assert!(matches!(
            st.verify_proof(&proof, &point_args, &message),
            Err(VerificationFailure)
        ));

        // Test bad proof #3 - incorrect # of scalars (1 too few)
        let mut proof2 = proof.clone();
        proof2.truncate(proof2.len() - 32);
        assert!(matches!(
            st.verify_proof(&proof2, &point_args, &message),
            Err(VerificationFailure)
        ));

        // Test bad proof #3 - incorrect # of scalars (1 too few)
        let mut proof2 = proof.clone();
        proof2.truncate(proof2.len() - 32);
        assert!(matches!(
            st.verify_proof(&proof2, &point_args, &message),
            Err(VerificationFailure)
        ));
    }
}
