use crate::{
    Ciphersuite, Element, Scalar,
    crypto::{
        constants::{
            NEAR_DLOG_CHALLENGE_LABEL, NEAR_DLOG_COMMITMENT_LABEL, NEAR_DLOG_ENCODE_LABEL_PUBLIC,
            NEAR_DLOG_ENCODE_LABEL_STATEMENT, NEAR_DLOG_STATEMENT_LABEL,
        },
        proofs::strobe_transcript::TranscriptRng,
    },
    errors::ProtocolError,
};
use frost_core::{Group, serialization::SerializableScalar};
use subtle::ConstantTimeEq;
use zeroize::Zeroize;

use super::strobe_transcript::Transcript;

/// The public statement for this proof.
/// This statement claims knowledge of the discrete logarithm of some point.
#[derive(Clone, Copy)]
pub struct Statement<'a, C: Ciphersuite> {
    pub public: &'a Element<C>,
}

impl<C: Ciphersuite> Statement<'_, C> {
    /// Encode into Vec<u8>: some sort of serialization
    fn encode(self) -> Result<Vec<u8>, ProtocolError> {
        let mut enc = Vec::new();
        enc.extend_from_slice(NEAR_DLOG_ENCODE_LABEL_STATEMENT);

        match <C::Group as Group>::serialize(self.public) {
            Ok(ser) => {
                enc.extend_from_slice(NEAR_DLOG_ENCODE_LABEL_PUBLIC);
                enc.extend_from_slice(ser.as_ref());
            }
            _ => return Err(ProtocolError::PointSerialization),
        }
        Ok(enc)
    }
}

/// The private witness for this proof.
/// This holds the scalar the prover needs to know.
#[derive(Clone)]
pub struct Witness<C: Ciphersuite>
where
    crate::Scalar<C>: Zeroize,
{
    pub x: SerializableScalar<C>,
}

impl<C: Ciphersuite> Zeroize for Witness<C>
where
    crate::Scalar<C>: Zeroize,
{
    fn zeroize(&mut self) {
        self.x.0.zeroize();
    }
}

impl<C: Ciphersuite> Drop for Witness<C>
where
    crate::Scalar<C>: Zeroize,
{
    fn drop(&mut self) {
        self.zeroize();
    }
}

/// Represents a proof of the statement.
#[derive(Clone, serde::Serialize, serde::Deserialize)]
#[serde(bound = "C: Ciphersuite")]
pub struct Proof<C: Ciphersuite> {
    e: SerializableScalar<C>,
    s: SerializableScalar<C>,
}

/// Produce a proof for the given statement and witness, using a caller-provided nonce.
/// The nonce `(k, k * G)` must be sampled from a cryptographically secure RNG by the caller.
/// The challenge is derived via the Fiat-Shamir transform over the transcript.
pub fn prove_with_nonce<C: Ciphersuite>(
    transcript: &mut Transcript,
    statement: Statement<'_, C>,
    witness: &Witness<C>,
    nonce: (Scalar<C>, Element<C>),
) -> Result<Proof<C>, ProtocolError>
where
    Scalar<C>: Zeroize,
{
    transcript.message(NEAR_DLOG_STATEMENT_LABEL, &statement.encode()?);

    let (k, big_k) = nonce;

    // Create a serialization of big_k
    let ser = C::Group::serialize(&big_k).map_err(|_| ProtocolError::IdentityElement)?;
    transcript.message(NEAR_DLOG_COMMITMENT_LABEL, ser.as_ref());
    let mut rng = transcript.challenge_then_build_rng(NEAR_DLOG_CHALLENGE_LABEL);
    let e = frost_core::random_nonzero::<C, _>(&mut rng);

    let s = k + e * witness.x.0;
    Ok(Proof {
        e: SerializableScalar(e),
        s: SerializableScalar(s),
    })
}

/// Verify that a proof attesting to the validity of some statement.
/// We use a transcript in order to verify the Fiat-Shamir transformation.
pub fn verify<C: Ciphersuite>(
    transcript: &mut Transcript,
    statement: Statement<'_, C>,
    proof: &Proof<C>,
) -> Result<bool, ProtocolError>
where
    Scalar<C>: ConstantTimeEq,
{
    transcript.message(NEAR_DLOG_STATEMENT_LABEL, &statement.encode()?);

    let big_k = C::Group::generator() * proof.s.0 - *statement.public * proof.e.0;

    // Create a serialization of big_k
    // Raises error if the big_k turned out to be the identity element
    let ser = C::Group::serialize(&big_k).map_err(|_| ProtocolError::IdentityElement)?;

    transcript.message(NEAR_DLOG_COMMITMENT_LABEL, ser.as_ref());
    let mut rng = transcript.challenge_then_build_rng(NEAR_DLOG_CHALLENGE_LABEL);
    let e = frost_core::random_nonzero::<C, TranscriptRng>(&mut rng);

    Ok(bool::from(e.ct_eq(&proof.e.0)))
}

#[cfg(test)]
mod test {
    use elliptic_curve::{bigint::Uint, scalar::FromUintUnchecked};

    use super::*;
    use frost_secp256k1::Secp256K1Sha256;
    use k256::{ProjectivePoint, Scalar};

    #[test]
    fn test_verify_fixed_randomness() {
        let x = Scalar::from_uint_unchecked(Uint::from_be_hex(
            "FC9A011DF3753BD79D841C11F6521F25AD2AB1DECEB96B7E8C28D87EA3303A06",
        ));
        let transcript = Transcript::new(b"protocol");
        let statement = Statement::<Secp256K1Sha256> {
            public: &(ProjectivePoint::GENERATOR * x),
        };
        let proof: Proof<Secp256K1Sha256> = Proof {
            e: SerializableScalar(Scalar::from_uint_unchecked(Uint::from_be_hex(
                "BA7718DDF60BC62FC6081B658322E908CD4FF161AB754748EC170CBC66898CDB",
            ))),
            s: SerializableScalar(Scalar::from_uint_unchecked(Uint::from_be_hex(
                "5086B275DC32C8CD1AAD377918E0B622BAF92844BDC46808BD5568D6E304DB33",
            ))),
        };
        assert!(verify(&mut transcript.fork(b"party", &[1]), statement, &proof).unwrap());
    }
}