use crate::cryptography::{Ciphertext, HybridCiphertext, PublicKey, SecretKey};
use crate::encrypted_vote::{EncryptedVote, ProofOfCorrectVote, Vote};
use crate::math::polynomial::Polynomial;
use crate::tally::Crs;
use crate::{GroupElement, Scalar, CURVE_HRP};
use chain_crypto::bech32::{to_bech32_from_bytes, try_from_bech32_to_bytes, Bech32, Error};
use const_format::concatcp;
use rand_core::{CryptoRng, RngCore};

/// Committee member election secret key
#[derive(Clone)]
pub struct MemberSecretKey(pub(crate) SecretKey);

/// Committee member election public key
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct MemberPublicKey(pub(crate) PublicKey);

#[derive(Clone)]
pub struct MemberCommunicationKey(SecretKey);

/// Committee Member communication public key (with other committee members)
#[derive(Clone)]
pub struct MemberCommunicationPublicKey(PublicKey);

/// The overall committee public key used for everyone to encrypt their vote to.
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct ElectionPublicKey(pub(crate) PublicKey);

impl ElectionPublicKey {
    #[doc(hidden)]
    pub fn as_raw(&self) -> &PublicKey {
        &self.0
    }

    /// Take a vote and encrypt it + provide a proof of correct voting
    pub fn encrypt_and_prove_vote<R: RngCore + CryptoRng>(
        &self,
        rng: &mut R,
        crs: &Crs,
        vote: Vote,
    ) -> (EncryptedVote, ProofOfCorrectVote) {
        let encryption_randomness = vec![Scalar::random(rng); vote.len()];
        let ciphertexts: Vec<Ciphertext> = encryption_randomness
            .iter()
            .zip(vote.iter())
            .map(|(r, v)| self.as_raw().encrypt_with_r(&Scalar::from(v), r))
            .collect();

        let proof = ProofOfCorrectVote::generate(
            rng,
            crs,
            &self.0,
            &vote,
            &encryption_randomness,
            &ciphertexts,
        );
        (ciphertexts, proof)
    }

    /// Create an election public key from all the participants of this committee
    pub fn from_participants(pks: &[MemberPublicKey]) -> Self {
        let mut k = pks[0].0.pk.clone();
        for pk in &pks[1..] {
            k = k + &pk.0.pk;
        }
        ElectionPublicKey(PublicKey { pk: k })
    }

    pub fn to_bytes(&self) -> Vec<u8> {
        self.0.to_bytes()
    }

    pub fn from_bytes(buf: &[u8]) -> Option<Self> {
        PublicKey::from_bytes(buf).map(ElectionPublicKey)
    }
}

impl Bech32 for ElectionPublicKey {
    const BECH32_HRP: &'static str = concatcp!(CURVE_HRP, "_votepk");
    const BYTES_LEN: usize = PublicKey::BYTES_LEN;

    fn try_from_bech32_str(bech32_str: &str) -> Result<Self, Error> {
        try_from_bech32_to_bytes::<Self>(bech32_str).and_then(|raw| {
            Self::from_bytes(&raw)
                .ok_or_else(|| Error::DataInvalid("invalid election public key binary data".into()))
        })
    }

    fn to_bech32_str(&self) -> String {
        to_bech32_from_bytes::<Self>(&self.to_bytes())
    }
}

/// Initial state generated by a Member, which include keys for this election
#[derive(Clone)]
// TODO: some of the fields are needed for DKG features which are
// not yet implemented
#[allow(dead_code)]
pub struct MemberState {
    sk: MemberSecretKey,
    owner_index: usize,
    apubs: Vec<GroupElement>,
    es: Vec<GroupElement>,
    encrypted: Vec<(HybridCiphertext, HybridCiphertext)>,
}

impl MemberState {
    /// Generate a new member state from random, where the number
    pub fn new<R: RngCore + CryptoRng>(
        rng: &mut R,
        t: usize,
        h: &Crs, // TODO: document
        committee_pks: &[MemberCommunicationPublicKey],
        my: usize,
    ) -> MemberState {
        let n = committee_pks.len();
        assert!(t > 0);
        assert!(t <= n);
        assert!(my < n);

        let pcomm = Polynomial::random(rng, t);
        let pshek = Polynomial::random(rng, t);

        let mut apubs = Vec::new();
        let mut es = Vec::new();

        for (ai, bi) in pshek.get_coefficients().zip(pcomm.get_coefficients()) {
            let apub = GroupElement::generator() * ai;
            let e = &apub + h * bi;
            apubs.push(apub);
            es.push(e);
        }

        let mut encrypted = Vec::new();
        #[allow(clippy::needless_range_loop)]
        for i in 0..n {
            // don't generate share for self
            if i == my {
                continue;
            } else {
                let idx = Scalar::from_u64((i + 1) as u64);
                let share_comm = pcomm.evaluate(&idx);
                let share_shek = pshek.evaluate(&idx);

                let pk = &committee_pks[i];

                let ecomm = pk.0.hybrid_encrypt(&share_comm.to_bytes(), rng);
                let eshek = pk.0.hybrid_encrypt(&share_shek.to_bytes(), rng);

                encrypted.push((ecomm, eshek));
            }
        }

        assert_eq!(apubs.len(), t + 1);
        assert_eq!(es.len(), t + 1);
        assert_eq!(encrypted.len(), n - 1);

        MemberState {
            sk: MemberSecretKey(SecretKey {
                sk: pshek.at_zero(),
            }),
            owner_index: my + 1, // committee member are 1-indexed
            apubs,
            es,
            encrypted,
        }
    }

    pub fn secret_key(&self) -> &MemberSecretKey {
        &self.sk
    }

    pub fn member_secret_key(&self) -> MemberSecretKey {
        self.sk.clone()
    }

    pub fn public_key(&self) -> MemberPublicKey {
        MemberPublicKey(PublicKey {
            pk: self.apubs[0].clone(),
        })
    }
}

impl MemberSecretKey {
    pub fn to_bytes(&self) -> [u8; 32] {
        self.0.sk.to_bytes()
    }

    pub fn from_bytes(bytes: &[u8]) -> Option<Self> {
        let sk = Scalar::from_bytes(bytes)?;
        Some(Self(SecretKey { sk }))
    }

    pub fn to_public(&self) -> MemberPublicKey {
        MemberPublicKey(PublicKey {
            pk: GroupElement::generator() * &self.0.sk,
        })
    }
}

impl Bech32 for MemberSecretKey {
    const BECH32_HRP: &'static str = concatcp!(CURVE_HRP, "_membersk");
    const BYTES_LEN: usize = SecretKey::BYTES_LEN;

    fn try_from_bech32_str(bech32_str: &str) -> Result<Self, Error> {
        try_from_bech32_to_bytes::<Self>(bech32_str).and_then(|raw| {
            Self::from_bytes(&raw)
                .ok_or_else(|| Error::DataInvalid("invalid member secret key binary data".into()))
        })
    }

    fn to_bech32_str(&self) -> String {
        to_bech32_from_bytes::<Self>(&self.to_bytes())
    }
}

impl MemberPublicKey {
    pub const BYTES_LEN: usize = PublicKey::BYTES_LEN;

    pub fn to_bytes(&self) -> Vec<u8> {
        self.0.to_bytes()
    }

    pub fn from_bytes(buf: &[u8]) -> Option<Self> {
        Some(Self(PublicKey::from_bytes(buf)?))
    }
}

impl Bech32 for MemberPublicKey {
    const BECH32_HRP: &'static str = concatcp!(CURVE_HRP, "_memberpk");
    const BYTES_LEN: usize = PublicKey::BYTES_LEN;

    fn try_from_bech32_str(bech32_str: &str) -> Result<Self, Error> {
        try_from_bech32_to_bytes::<Self>(bech32_str).and_then(|raw| {
            Self::from_bytes(&raw)
                .ok_or_else(|| Error::DataInvalid("invalid member public key binary data".into()))
        })
    }

    fn to_bech32_str(&self) -> String {
        to_bech32_from_bytes::<Self>(&self.to_bytes())
    }
}

impl From<PublicKey> for MemberPublicKey {
    fn from(pk: PublicKey) -> MemberPublicKey {
        MemberPublicKey(pk)
    }
}

impl MemberCommunicationKey {
    pub fn new<R: RngCore + CryptoRng>(rng: &mut R) -> Self {
        let sk = SecretKey::generate(rng);
        MemberCommunicationKey(sk)
    }

    pub fn to_public(&self) -> MemberCommunicationPublicKey {
        MemberCommunicationPublicKey(PublicKey {
            pk: &GroupElement::generator() * &self.0.sk,
        })
    }

    pub fn from_bytes(bytes: &[u8]) -> Option<MemberCommunicationKey> {
        SecretKey::from_bytes(bytes).map(MemberCommunicationKey)
    }

    pub fn to_bytes(&self) -> [u8; 32] {
        self.0.sk.to_bytes()
    }
}

impl Bech32 for MemberCommunicationKey {
    const BECH32_HRP: &'static str = concatcp!(CURVE_HRP, "_vcommsk");
    const BYTES_LEN: usize = SecretKey::BYTES_LEN;

    fn try_from_bech32_str(bech32_str: &str) -> Result<Self, Error> {
        try_from_bech32_to_bytes::<Self>(bech32_str).and_then(|raw| {
            Self::from_bytes(&raw).ok_or_else(|| {
                Error::DataInvalid("invalid member communication secret key binary data".into())
            })
        })
    }

    fn to_bech32_str(&self) -> String {
        to_bech32_from_bytes::<Self>(&self.to_bytes())
    }
}

impl From<PublicKey> for MemberCommunicationPublicKey {
    fn from(pk: PublicKey) -> MemberCommunicationPublicKey {
        Self(pk)
    }
}

impl MemberCommunicationPublicKey {
    pub fn to_bytes(&self) -> Vec<u8> {
        self.0.to_bytes()
    }

    pub fn from_bytes(bytes: &[u8]) -> Option<Self> {
        PublicKey::from_bytes(bytes).map(Self)
    }
}

impl Bech32 for MemberCommunicationPublicKey {
    const BECH32_HRP: &'static str = concatcp!(CURVE_HRP, "_vcommpk");
    const BYTES_LEN: usize = PublicKey::BYTES_LEN;

    fn try_from_bech32_str(bech32_str: &str) -> Result<Self, Error> {
        try_from_bech32_to_bytes::<Self>(bech32_str).and_then(|raw| {
            Self::from_bytes(&raw).ok_or_else(|| {
                Error::DataInvalid("invalid member communication public key binary data".into())
            })
        })
    }

    fn to_bech32_str(&self) -> String {
        to_bech32_from_bytes::<Self>(&self.to_bytes())
    }
}