use crate::{
    core::{
        circuits::boolean::{boolean_value::BooleanValue, byte::Byte},
        global_value::value::FieldValue,
    },
    utils::{
        field::ScalarField,
        zkp::{
            elgamal::{DecryptHandle, ElGamalCiphertext, ElGamalPubkey},
            pedersen::{Pedersen, PedersenCommitment, PedersenOpening},
        },
    },
};
use zk_elgamal_proof::encryption::{DECRYPT_HANDLE_LEN, PEDERSEN_COMMITMENT_LEN};

pub const GROUPED_ELGAMAL_CIPHERTEXT_2_HANDLES_LEN: usize =
    PEDERSEN_COMMITMENT_LEN + 2 * DECRYPT_HANDLE_LEN;

pub const GROUPED_ELGAMAL_CIPHERTEXT_3_HANDLES_LEN: usize =
    PEDERSEN_COMMITMENT_LEN + 3 * DECRYPT_HANDLE_LEN;

/// Algorithm handle for the grouped ElGamal encryption
pub struct GroupedElGamal<const N: usize>;
impl<const N: usize> GroupedElGamal<N> {
    /// Encrypts an amount under an array of ElGamal public keys using a specified Pedersen
    /// opening.
    pub fn encrypt_with(
        pubkeys: [&ElGamalPubkey; N],
        amount: FieldValue<ScalarField>,
        opening: &PedersenOpening,
    ) -> GroupedElGamalCiphertext<N> {
        let commitment = Pedersen::with(amount, opening);
        let handles = pubkeys.map(|handle| handle.decrypt_handle(opening));

        GroupedElGamalCiphertext {
            commitment,
            handles,
        }
    }

    /// Converts a grouped ElGamal ciphertext into a regular ElGamal ciphertext using the decrypt
    /// handle at a specified index.
    fn to_elgamal_ciphertext(
        grouped_ciphertext: &GroupedElGamalCiphertext<N>,
        index: usize,
    ) -> ElGamalCiphertext {
        assert!(index < N);
        ElGamalCiphertext {
            commitment: grouped_ciphertext.commitment,
            handle: grouped_ciphertext.handles[index],
        }
    }
}

/// A grouped ElGamal ciphertext.
///
/// The type is defined with a generic constant parameter that specifies the number of
/// decryption handles that the ciphertext holds.
#[derive(Clone, Copy)]
pub struct GroupedElGamalCiphertext<const N: usize> {
    pub commitment: PedersenCommitment,
    pub handles: [DecryptHandle; N],
}

impl<const N: usize> GroupedElGamalCiphertext<N> {
    /// Converts a grouped ElGamal ciphertext into a regular ElGamal ciphertext using the decrypt
    /// handle at a specified index.
    pub fn to_elgamal_ciphertext(&self, index: usize) -> ElGamalCiphertext {
        GroupedElGamal::to_elgamal_ciphertext(self, index)
    }
}

#[derive(Clone, Copy)]
pub struct GroupedElGamalCiphertext2Handles(pub GroupedElGamalCiphertext<2>);

impl GroupedElGamalCiphertext2Handles {
    pub fn to_bytes(&self) -> [Byte<BooleanValue>; GROUPED_ELGAMAL_CIPHERTEXT_2_HANDLES_LEN] {
        let mut bytes = [Byte::<BooleanValue>::from(0u8); GROUPED_ELGAMAL_CIPHERTEXT_2_HANDLES_LEN];
        bytes[..PEDERSEN_COMMITMENT_LEN]
            .copy_from_slice(&self.0.commitment.get_point().compress().to_bytes());
        let mut offset = PEDERSEN_COMMITMENT_LEN;
        for handle in &self.0.handles {
            bytes[offset..offset + DECRYPT_HANDLE_LEN]
                .copy_from_slice(&handle.get_point().compress().to_bytes());
            offset += DECRYPT_HANDLE_LEN;
        }
        bytes
    }
}

#[derive(Clone, Copy)]
pub struct GroupedElGamalCiphertext3Handles(pub GroupedElGamalCiphertext<3>);

impl GroupedElGamalCiphertext3Handles {
    pub fn to_bytes(&self) -> [Byte<BooleanValue>; GROUPED_ELGAMAL_CIPHERTEXT_3_HANDLES_LEN] {
        let mut bytes = [Byte::<BooleanValue>::from(0u8); GROUPED_ELGAMAL_CIPHERTEXT_3_HANDLES_LEN];
        bytes[..PEDERSEN_COMMITMENT_LEN]
            .copy_from_slice(&self.0.commitment.get_point().compress().to_bytes());
        let mut offset = PEDERSEN_COMMITMENT_LEN;
        for handle in &self.0.handles {
            bytes[offset..offset + DECRYPT_HANDLE_LEN]
                .copy_from_slice(&handle.get_point().compress().to_bytes());
            offset += DECRYPT_HANDLE_LEN;
        }
        bytes
    }
}
