use std::borrow::Borrow;

use bincode::{deserialize, serialize};
use curve25519_dalek::ristretto::RistrettoPoint;
use curve25519_dalek::traits::{Identity, MultiscalarMul};
use curve25519_dalek::{constants, scalar};
use ring::{error, rand};
use serde::{Deserialize, Serialize};
use sha3::{Digest, Sha3_256};

use super::*;

#[derive(Debug, Serialize, Deserialize)]
struct SerializableAbeRingSignature {
    ring: Vec<([u8; 32], [u8; 32])>,
    s: [u8; 32],
}

#[derive(Clone)]
pub struct AbeRingSignature {
    ring: Vec<(PublicKey, scalar::Scalar)>,
    s: scalar::Scalar,
}

pub struct AbeRingSigner<D> {
    c: D,
    z: RistrettoPoint,
    ring: Vec<(PublicKey, scalar::Scalar)>,
}

impl AbeRingSignature {
    pub fn verify(&self, data: &[u8]) -> bool {
        self.verify_with_digest::<Sha3_256>(data)
    }

    // XXX: use Keccak TupleHash
    pub fn verify_with_digest<D: Digest>(&self, data: &[u8]) -> bool {
        let mut sum = scalar::Scalar::zero();
        let mut digest = D::new();

        let scalars = self.ring.iter().map(|&(_, ref c)| c);
        let points = self.ring.iter().map(|&(ref p, _)| &p.inner);
        // XXX: In theory, I don't think we actually *need* this multiscalar multiplication in
        // constant time.
        // It saves around 15% of time, for only a ring of two.
        let z = &self.s * &constants::RISTRETTO_BASEPOINT_TABLE
            + RistrettoPoint::multiscalar_mul(scalars, points);

        // TODO: implement Sumcurve25519_dalek::scalar::Scalar: std::iter::Sum<&curve25519_dalek::scalar::Scalar>
        // let sum: scalar::Scalar = scalars.sum();
        for &(ref member, ref c) in self.ring.iter() {
            sum += c;
            digest.update(member.inner.compress().as_bytes());
        }
        digest.update(data);
        digest.update(z.compress().as_bytes());

        let hash = scalar_from_digest(digest);

        sum -= hash;

        sum == scalar::Scalar::zero()
    }

    pub fn serialize(&self) -> Vec<u8> {
        serialize(&SerializableAbeRingSignature {
            ring: self
                .ring
                .iter()
                .map(|&(ref pk, ref scalar)| (pk.bytes(), scalar.to_bytes()))
                .collect(),
            s: self.s.to_bytes(),
        })
        .expect("Could not serialize RingSignature")
    }

    pub fn deserialize<B: AsRef<[u8]>>(bytes: B) -> Result<AbeRingSignature, error::Unspecified> {
        let signature = deserialize::<SerializableAbeRingSignature>(bytes.as_ref())
            .map_err(|_| error::Unspecified)?;

        Ok(AbeRingSignature {
            ring: signature
                .ring
                .into_iter()
                .map(|(pk, scalar)| {
                    (
                        PublicKey::from_bytes(&pk).expect("Invalid public key"),
                        scalar::Scalar::from_bytes_mod_order(scalar),
                    )
                })
                .collect(),
            s: scalar::Scalar::from_bytes_mod_order(signature.s),
        })
    }
}

/// Uses the Appendix A algorithm of the Abe et al. 2002.
impl<D: Digest> AbeRingSigner<D> {
    pub fn push<R: rand::SecureRandom>(
        &mut self,
        pk: &PublicKey,
        rng: &R,
    ) -> crate::crypto::Result<()> {
        let c = random_scalar(rng)?;
        // z += c_i * Y_i
        self.z += c * pk.inner;
        self.ring.push((pk.clone(), c));
        Ok(())
    }

    pub fn push_many<R, Key, KeyList>(&mut self, pk: KeyList, rng: &R) -> crate::crypto::Result<()>
    where
        KeyList: AsRef<[Key]>,
        R: rand::SecureRandom,
        Key: Borrow<PublicKey>,
    {
        let pk = pk.as_ref();
        let pk_iter = pk.iter().map(|k| k.borrow());

        let scalars: Result<Vec<_>, _> = (0..pk.len()).map(|_| random_scalar(rng)).collect();
        let scalars = scalars?;

        self.z += RistrettoPoint::multiscalar_mul(&scalars, pk_iter.clone().map(|key| key.inner));

        self.ring.extend(
            pk_iter
                .zip(scalars)
                .map(|(key, scalar)| (key.clone(), scalar)),
        );

        Ok(())
    }

    pub fn finish<R: rand::SecureRandom>(
        self,
        sk: &PrivateKey,
        position: usize,
        rng: &R,
        data: &[u8],
    ) -> crate::crypto::Result<AbeRingSignature> {
        let alpha = random_scalar(rng)?;
        let AbeRingSigner {
            mut c,
            mut z,
            mut ring,
        } = self;

        // Finish z
        z += &alpha * &constants::RISTRETTO_BASEPOINT_TABLE;

        // Finish c
        for item in &ring[0..position] {
            c.update(item.0.inner.compress().as_bytes());
        }
        c.update(sk.compute_public_key()?.inner.compress().as_bytes());
        for item in &ring[position..ring.len()] {
            c.update(item.0.inner.compress().as_bytes());
        }
        c.update(data);
        c.update(z.compress().as_bytes());

        let mut ck = scalar_from_digest(c);

        let mut sum = scalar::Scalar::zero();
        for &(_, ref c) in ring.iter() {
            sum += c;
        }
        ck -= sum;
        ring.insert(position, (sk.compute_public_key()?, ck));

        // calculate s = alpha - c_k*sk
        let s = alpha - ck * sk.inner;

        let signature = AbeRingSignature { ring, s };

        Ok(signature)
    }
}

impl Default for AbeRingSigner<Sha3_256> {
    fn default() -> Self {
        AbeRingSigner {
            c: Sha3_256::default(),
            z: RistrettoPoint::identity(),
            ring: Vec::new(),
        }
    }
}