use std::io::{self, Read, Write};

use zeroize::Zeroize;

use transcript::{Transcript, SecureDigest, DigestTranscript};

use group::{
  ff::{Field, PrimeField},
  Group, GroupEncoding,
};
use multiexp::multiexp_vartime;
use ciphersuite::Ciphersuite;

use crate::SchnorrSignature;

// Performs a big-endian modular reduction of the hash value
// This is used by the below aggregator to prevent mutability
// Only an 128-bit scalar is needed to offer 128-bits of security against malleability per
// https://cr.yp.to/badbatch/badbatch-20120919.pdf
// Accordingly, while a 256-bit hash used here with a 256-bit ECC will have bias, it shouldn't be
// an issue
fn scalar_from_digest<D: Clone + SecureDigest, F: PrimeField>(
  digest: &mut DigestTranscript<D>,
) -> F {
  let bytes = digest.challenge(b"aggregation_weight");
  debug_assert_eq!(bytes.len() % 8, 0);

  let mut res = F::zero();
  let mut i = 0;
  while i < bytes.len() {
    if i != 0 {
      for _ in 0 .. 8 {
        res += res;
      }
    }
    res += F::from(u64::from_be_bytes(bytes[i .. (i + 8)].try_into().unwrap()));
    i += 8;
  }
  res
}

/// Aggregate Schnorr signature as defined in <https://eprint.iacr.org/2021/350>.
#[allow(non_snake_case)]
#[derive(Clone, PartialEq, Eq, Debug, Zeroize)]
pub struct SchnorrAggregate<C: Ciphersuite> {
  pub Rs: Vec<C::G>,
  pub s: C::F,
}

impl<C: Ciphersuite> SchnorrAggregate<C> {
  /// Read a SchnorrAggregate from something implementing Read.
  pub fn read<R: Read>(reader: &mut R) -> io::Result<Self> {
    let mut len = [0; 4];
    reader.read_exact(&mut len)?;

    #[allow(non_snake_case)]
    let mut Rs = vec![];
    for _ in 0 .. u32::from_le_bytes(len) {
      Rs.push(C::read_G(reader)?);
    }

    Ok(SchnorrAggregate { Rs, s: C::read_F(reader)? })
  }

  /// Write a SchnorrAggregate to something implementing Write.
  pub fn write<W: Write>(&self, writer: &mut W) -> io::Result<()> {
    writer.write_all(
      &u32::try_from(self.Rs.len())
        .expect("more than 4 billion signatures in aggregate")
        .to_le_bytes(),
    )?;
    #[allow(non_snake_case)]
    for R in &self.Rs {
      writer.write_all(R.to_bytes().as_ref())?;
    }
    writer.write_all(self.s.to_repr().as_ref())
  }

  /// Serialize a SchnorrAggregate, returning a Vec<u8>.
  pub fn serialize(&self) -> Vec<u8> {
    let mut buf = vec![];
    self.write(&mut buf).unwrap();
    buf
  }

  /// Perform signature verification.
  ///
  /// Challenges must be properly crafted, which means being binding to the public key, nonce, and
  /// any message. Failure to do so will let a malicious adversary to forge signatures for
  /// different keys/messages.
  ///
  /// The DST used here must prevent a collision with whatever hash function produced the
  /// challenges.
  #[must_use]
  pub fn verify<D: Clone + SecureDigest>(
    &self,
    dst: &'static [u8],
    keys_and_challenges: &[(C::G, C::F)],
  ) -> bool {
    if self.Rs.len() != keys_and_challenges.len() {
      return false;
    }

    let mut digest = DigestTranscript::<D>::new(dst);
    digest.domain_separate(b"signatures");
    for (_, challenge) in keys_and_challenges {
      digest.append_message(b"challenge", challenge.to_repr());
    }

    let mut pairs = Vec::with_capacity((2 * keys_and_challenges.len()) + 1);
    for (i, (key, challenge)) in keys_and_challenges.iter().enumerate() {
      let z = scalar_from_digest(&mut digest);
      pairs.push((z, self.Rs[i]));
      pairs.push((z * challenge, *key));
    }
    pairs.push((-self.s, C::generator()));
    multiexp_vartime(&pairs).is_identity().into()
  }
}

#[allow(non_snake_case)]
#[derive(Clone, Debug, Zeroize)]
pub struct SchnorrAggregator<D: Clone + SecureDigest, C: Ciphersuite> {
  digest: DigestTranscript<D>,
  sigs: Vec<SchnorrSignature<C>>,
}

impl<D: Clone + SecureDigest, C: Ciphersuite> SchnorrAggregator<D, C> {
  /// Create a new aggregator.
  ///
  /// The DST used here must prevent a collision with whatever hash function produced the
  /// challenges.
  pub fn new(dst: &'static [u8]) -> Self {
    let mut res = Self { digest: DigestTranscript::<D>::new(dst), sigs: vec![] };
    res.digest.domain_separate(b"signatures");
    res
  }

  /// Aggregate a signature.
  pub fn aggregate(&mut self, challenge: C::F, sig: SchnorrSignature<C>) {
    self.digest.append_message(b"challenge", challenge.to_repr());
    self.sigs.push(sig);
  }

  /// Complete aggregation, returning None if none were aggregated.
  pub fn complete(mut self) -> Option<SchnorrAggregate<C>> {
    if self.sigs.is_empty() {
      return None;
    }

    let mut aggregate =
      SchnorrAggregate { Rs: Vec::with_capacity(self.sigs.len()), s: C::F::zero() };
    for i in 0 .. self.sigs.len() {
      aggregate.Rs.push(self.sigs[i].R);
      aggregate.s += self.sigs[i].s * scalar_from_digest::<_, C::F>(&mut self.digest);
    }
    Some(aggregate)
  }
}