//! Arcis implementation of <https://github.com/dalek-cryptography/merlin/blob/master/src/transcript.rs>

use crate::{
    core::{
        circuits::boolean::{
            boolean_value::{Boolean, BooleanValue},
            byte::Byte,
        },
        global_value::{curve_value::CompressedCurveValue, value::FieldValue},
    },
    traits::{FromLeBits, ToLeBytes},
    utils::{
        field::ScalarField,
        zkp::{
            elgamal::ElGamalCiphertext,
            grouped_elgamal::{
                GroupedElGamalCiphertext2Handles,
                GroupedElGamalCiphertext3Handles,
                GROUPED_ELGAMAL_CIPHERTEXT_2_HANDLES_LEN,
                GROUPED_ELGAMAL_CIPHERTEXT_3_HANDLES_LEN,
            },
            strobe::Strobe128,
        },
    },
};
use zk_elgamal_proof::encryption::{
    DECRYPT_HANDLE_LEN,
    ELGAMAL_CIPHERTEXT_LEN,
    PEDERSEN_COMMITMENT_LEN,
};

/// Domain separation label to initialize the STROBE context.
///
/// This is not to be confused with the crate's semver string:
/// the latter applies to the API, while this label defines the protocol.
/// E.g. it is possible that crate 2.0 will have an incompatible API,
/// but implement the same 1.0 protocol.
pub const MERLIN_PROTOCOL_LABEL: &[u8] = b"Merlin v1.0";

fn encode_u64(x: u64) -> [u8; 8] {
    use byteorder::{ByteOrder, LittleEndian};

    let mut buf = [0; 8];
    LittleEndian::write_u64(&mut buf, x);
    buf
}

fn encode_usize_as_u32(x: usize) -> [u8; 4] {
    use byteorder::{ByteOrder, LittleEndian};

    assert!(x <= (u32::MAX as usize));

    let mut buf = [0; 4];
    LittleEndian::write_u32(&mut buf, x as u32);
    buf
}

/// A transcript of a public-coin argument.
///
/// The prover's messages are added to the transcript using
/// [`append_message`](Transcript::append_message), and the verifier's
/// challenges can be computed using
/// [`challenge_bytes`](Transcript::challenge_bytes).
///
/// # Creating and using a Merlin transcript
///
/// To create a Merlin transcript, use [`Transcript::new()`].  This
/// function takes a domain separation label which should be unique to
/// the application.
///
/// To use the transcript with a Merlin-based proof implementation,
/// the prover's side creates a Merlin transcript with an
/// application-specific domain separation label, and passes a `&mut`
/// reference to the transcript to the proving function(s).
///
/// To verify the resulting proof, the verifier creates their own
/// Merlin transcript using the same domain separation label, then
/// passes a `&mut` reference to the verifier's transcript to the
/// verification function.
///
/// # Implementing proofs using Merlin
///
/// For information on the design of Merlin and how to use it to
/// implement a proof system, see the documentation at
/// [merlin.cool](https://merlin.cool), particularly the [Using
/// Merlin](https://merlin.cool/use/index.html) section.
#[derive(Clone)]
pub struct Transcript<B: Boolean> {
    strobe: Strobe128<B>,
}

impl<B: Boolean> Transcript<B> {
    /// Initialize a new transcript with the supplied `label`, which
    /// is used as a domain separator.
    ///
    /// # Note
    ///
    /// This function should be called by a proof library's API
    /// consumer (i.e., the application using the proof library), and
    /// **not by the proof implementation**.  See the [Passing
    /// Transcripts](https://merlin.cool/use/passing.html) section of
    /// the Merlin website for more details on why.
    pub fn new(label: &'static [u8]) -> Transcript<B> {
        let mut transcript = Transcript {
            strobe: Strobe128::new(MERLIN_PROTOCOL_LABEL),
        };
        transcript.append_message(
            b"dom-sep",
            &label
                .iter()
                .copied()
                .map(Byte::<B>::from)
                .collect::<Vec<Byte<B>>>(),
        );

        transcript
    }

    /// Append a prover's `message` to the transcript.
    ///
    /// The `label` parameter is metadata about the message, and is
    /// also appended to the transcript.  See the [Transcript
    /// Protocols](https://merlin.cool/use/protocol.html) section of
    /// the Merlin website for details on labels.
    pub fn append_message(&mut self, label: &'static [u8], message: &[Byte<B>]) {
        let data_len = encode_usize_as_u32(message.len());
        self.strobe.meta_ad(label, false);
        self.strobe.meta_ad(&data_len, true);
        self.strobe.ad(message, false);
    }

    /// Convenience method for appending a `u64` to the transcript.
    ///
    /// The `label` parameter is metadata about the message, and is
    /// also appended to the transcript.  See the [Transcript
    /// Protocols](https://merlin.cool/use/protocol.html) section of
    /// the Merlin website for details on labels.
    ///
    /// # Implementation
    ///
    /// Calls `append_message` with the 8-byte little-endian encoding
    /// of `x`.
    fn append_u64(&mut self, label: &'static [u8], x: u64) {
        self.append_message(
            label,
            &encode_u64(x)
                .into_iter()
                .map(Byte::<B>::from)
                .collect::<Vec<Byte<B>>>(),
        );
    }

    /// Fill the supplied buffer with the verifier's challenge bytes.
    ///
    /// The `label` parameter is metadata about the challenge, and is
    /// also appended to the transcript.  See the [Transcript
    /// Protocols](https://merlin.cool/use/protocol.html) section of
    /// the Merlin website for details on labels.
    pub fn challenge_bytes(&mut self, label: &'static [u8], dest: &mut [Byte<B>]) {
        let data_len = encode_usize_as_u32(dest.len());
        self.strobe.meta_ad(label, false);
        self.strobe.meta_ad(&data_len, true);
        self.strobe.prf(dest, false);
    }

    pub fn range_proof_domain_separator(&mut self, n: u64) {
        self.append_message(
            b"dom-sep",
            &b"range-proof"
                .iter()
                .copied()
                .map(Byte::<B>::from)
                .collect::<Vec<Byte<B>>>(),
        );
        self.append_u64(b"n", n);
    }

    pub fn inner_product_proof_domain_separator(&mut self, n: u64) {
        self.append_message(
            b"dom-sep",
            &b"inner-product"
                .iter()
                .copied()
                .map(Byte::<B>::from)
                .collect::<Vec<Byte<B>>>(),
        );
        self.append_u64(b"n", n);
    }

    pub fn ciphertext_commitment_equality_proof_domain_separator(&mut self) {
        self.append_message(
            b"dom-sep",
            &b"ciphertext-commitment-equality-proof"
                .iter()
                .copied()
                .map(Byte::<B>::from)
                .collect::<Vec<Byte<B>>>(),
        )
    }

    pub fn zero_ciphertext_proof_domain_separator(&mut self) {
        self.append_message(
            b"dom-sep",
            &b"zero-ciphertext-proof"
                .iter()
                .copied()
                .map(Byte::<B>::from)
                .collect::<Vec<Byte<B>>>(),
        )
    }

    pub fn grouped_ciphertext_validity_proof_domain_separator(&mut self, handles: u64) {
        self.append_message(
            b"dom-sep",
            &b"validity-proof"
                .iter()
                .copied()
                .map(Byte::<B>::from)
                .collect::<Vec<Byte<B>>>(),
        );
        self.append_u64(b"handles", handles);
    }

    pub fn batched_grouped_ciphertext_validity_proof_domain_separator(&mut self, handles: u64) {
        self.append_message(
            b"dom-sep",
            &b"batched-validity-proof"
                .iter()
                .copied()
                .map(Byte::<B>::from)
                .collect::<Vec<Byte<B>>>(),
        );
        self.append_u64(b"handles", handles);
    }

    pub fn percentage_with_cap_proof_domain_separator(&mut self) {
        self.append_message(
            b"dom-sep",
            &b"percentage-with-cap-proof"
                .iter()
                .copied()
                .map(Byte::<B>::from)
                .collect::<Vec<Byte<B>>>(),
        )
    }

    pub fn pubkey_proof_domain_separator(&mut self) {
        self.append_message(
            b"dom-sep",
            &b"pubkey-proof"
                .iter()
                .copied()
                .map(Byte::<B>::from)
                .collect::<Vec<Byte<B>>>(),
        )
    }
}

impl Transcript<BooleanValue> {
    pub fn append_scalar(&mut self, label: &'static [u8], scalar: &FieldValue<ScalarField>) {
        self.append_message(label, &scalar.to_le_bytes());
    }

    pub fn append_point(&mut self, label: &'static [u8], point: &CompressedCurveValue) {
        self.append_message(label, &point.to_bytes());
    }

    pub fn append_elgamal_ciphertext(
        &mut self,
        label: &'static [u8],
        ciphertext: &ElGamalCiphertext,
    ) {
        let mut bytes = [Byte::<BooleanValue>::from(0u8); ELGAMAL_CIPHERTEXT_LEN];
        bytes[..PEDERSEN_COMMITMENT_LEN]
            .copy_from_slice(&ciphertext.commitment.get_point().compress().to_bytes());
        bytes[PEDERSEN_COMMITMENT_LEN..]
            .copy_from_slice(&ciphertext.handle.get_point().compress().to_bytes());
        self.append_message(label, &bytes);
    }

    pub fn append_elgamal_ciphertext_2_handles(
        &mut self,
        label: &'static [u8],
        ciphertext: &GroupedElGamalCiphertext2Handles,
    ) {
        let mut bytes = [Byte::<BooleanValue>::from(0u8); GROUPED_ELGAMAL_CIPHERTEXT_2_HANDLES_LEN];
        bytes[..PEDERSEN_COMMITMENT_LEN]
            .copy_from_slice(&ciphertext.0.commitment.get_point().compress().to_bytes());
        let mut offset = PEDERSEN_COMMITMENT_LEN;
        for i in 0..2 {
            bytes[offset..offset + DECRYPT_HANDLE_LEN]
                .copy_from_slice(&ciphertext.0.handles[i].get_point().compress().to_bytes());
            offset += DECRYPT_HANDLE_LEN;
        }
        self.append_message(label, &bytes);
    }

    pub fn append_elgamal_ciphertext_3_handles(
        &mut self,
        label: &'static [u8],
        ciphertext: &GroupedElGamalCiphertext3Handles,
    ) {
        let mut bytes = [Byte::<BooleanValue>::from(0u8); GROUPED_ELGAMAL_CIPHERTEXT_3_HANDLES_LEN];
        bytes[..PEDERSEN_COMMITMENT_LEN]
            .copy_from_slice(&ciphertext.0.commitment.get_point().compress().to_bytes());
        let mut offset = PEDERSEN_COMMITMENT_LEN;
        for i in 0..3 {
            bytes[offset..offset + DECRYPT_HANDLE_LEN]
                .copy_from_slice(&ciphertext.0.handles[i].get_point().compress().to_bytes());
            offset += DECRYPT_HANDLE_LEN;
        }
        self.append_message(label, &bytes);
    }

    pub fn challenge_scalar(&mut self, label: &'static [u8]) -> FieldValue<ScalarField> {
        let mut buf = [Byte::from(0u8); 64];
        self.challenge_bytes(label, &mut buf);
        let buf_bits = buf
            .into_iter()
            .flat_map(|byte| byte.to_vec())
            .collect::<Vec<BooleanValue>>();

        FieldValue::<ScalarField>::from_le_bits(buf_bits, false)
    }
}
