// Copyright (c) Silence Laboratories Pte. Ltd. All Rights Reserved.
// This software is licensed under the Silence Laboratories License Agreement.

use crypto_box::{
    aead::{Aead, AeadCore},
    PublicKey, SalsaBox, SecretKey,
};

use crypto_bigint::{
    generic_array::{typenum::Unsigned, GenericArray},
    rand_core::CryptoRngCore,
};

use sha2::{Digest, Sha256};

// Encryption is done inplace, so the size of the ciphertext is the size of the message plus the tag size.
pub const SCALAR_CIPHERTEXT_SIZE: usize = 64 + <SalsaBox as AeadCore>::TagSize::USIZE;

// Custom serde serializer
#[cfg(feature = "serde")]
pub mod serde_point {
    use std::marker::PhantomData;

    use elliptic_curve::group::GroupEncoding;
    use serde::de::Visitor;

    pub fn serialize<S, G: GroupEncoding>(point: &G, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        use serde::ser::SerializeTuple;
        let mut tup = serializer.serialize_tuple(G::Repr::default().as_ref().len())?;
        for byte in point.to_bytes().as_ref().iter() {
            tup.serialize_element(byte)?;
        }
        tup.end()
    }

    pub fn deserialize<'de, D, G: GroupEncoding>(deserializer: D) -> Result<G, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        struct PointVisitor<G: GroupEncoding>(PhantomData<G>);

        impl<'de, G: GroupEncoding> Visitor<'de> for PointVisitor<G> {
            type Value = G;

            fn expecting(&self, formatter: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
                formatter.write_str("a valid point in Edwards y + sign format")
            }

            fn visit_seq<A>(self, mut seq: A) -> Result<G, A::Error>
            where
                A: serde::de::SeqAccess<'de>,
            {
                let mut encoding = G::Repr::default();
                for (idx, byte) in encoding.as_mut().iter_mut().enumerate() {
                    *byte = seq.next_element()?.ok_or_else(|| {
                        serde::de::Error::invalid_length(idx, &"wrong length of point")
                    })?;
                }

                Option::from(G::from_bytes(&encoding))
                    .ok_or(serde::de::Error::custom("point decompression failed"))
            }
        }

        deserializer.deserialize_tuple(G::Repr::default().as_ref().len(), PointVisitor(PhantomData))
    }
}

#[cfg(feature = "serde")]
pub mod serde_vec_point {
    use elliptic_curve::group::GroupEncoding;

    pub fn serialize<S, G: GroupEncoding>(points: &[G], serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        let mut bytes = Vec::with_capacity(points.len() * G::Repr::default().as_ref().len());
        points.iter().for_each(|point| {
            bytes.extend_from_slice(point.to_bytes().as_ref());
        });
        serializer.serialize_bytes(&bytes)
    }

    pub fn deserialize<'de, D, G: GroupEncoding>(deserializer: D) -> Result<Vec<G>, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let bytes: Vec<u8> = serde::Deserialize::deserialize(deserializer)?;
        let point_size = G::Repr::default().as_ref().len();
        if !bytes.len().is_multiple_of(point_size) {
            return Err(serde::de::Error::custom("Invalid number of bytes"));
        }
        let mut points = Vec::with_capacity(bytes.len() / point_size);
        for i in 0..bytes.len() / point_size {
            let mut encoding = G::Repr::default();
            encoding
                .as_mut()
                .copy_from_slice(&bytes[i * point_size..(i + 1) * point_size]);
            points.push(
                Option::from(G::from_bytes(&encoding))
                    .ok_or_else(|| serde::de::Error::custom("Invalid point"))?,
            );
        }
        Ok(points)
    }
}

#[cfg(feature = "serde")]
pub mod serde_arc {
    use std::sync::Arc;

    use serde::{Deserialize, Serialize};

    pub fn serialize<S, T: Serialize>(value: &Arc<T>, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        value.serialize(serializer)
    }

    pub fn deserialize<'de, D, T: Deserialize<'de>>(deserializer: D) -> Result<Arc<T>, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        Ok(Arc::new(T::deserialize(deserializer)?))
    }
}

#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[derive(Clone, Copy, Debug, bytemuck::AnyBitPattern, bytemuck::NoUninit)]
#[repr(C)]
/// Encrypted data
/// Specific to the keygen protocol
/// It is an encryption of 64 bytes of data.
/// The data:  [polynomial evaluation c_i_j (32 bytes) || chain_code_session_id (32 bytes)] for the party_{i}
pub struct EncryptedData {
    pub sender_pid: u8,
    pub receiver_pid: u8,
    #[cfg_attr(feature = "serde", serde(with = "serde_bytes"))]
    pub ciphertext: [u8; SCALAR_CIPHERTEXT_SIZE],
    // Size of the nonce is 24
    pub nonce: [u8; <SalsaBox as AeadCore>::NonceSize::USIZE],
}

// Not using the SessionId type from sl_mpc_mate because it is not serializable.
pub type SessionId = [u8; 32];
pub type HashBytes = [u8; 32];

impl EncryptedData {
    pub fn new(
        ciphertext: [u8; SCALAR_CIPHERTEXT_SIZE],
        nonce: [u8; 24],
        sender_pid: u8,
        receiver_pid: u8,
    ) -> Self {
        Self {
            ciphertext,
            nonce,
            sender_pid,
            receiver_pid,
        }
    }
}

/// Common trait for all MPC messages.
pub(crate) trait BaseMessage {
    // Return the party id of the message sender.
    fn party_id(&self) -> u8;
}

/// Calculates the final session id from the list of session ids.
pub fn calculate_final_session_id(
    party_ids: impl Iterator<Item = u8>,
    sid_i_list: &[SessionId],
    extra: &[&[u8]],
) -> SessionId {
    let mut hasher = Sha256::new();

    for e in extra {
        hasher.update(e);
    }

    party_ids.for_each(|pid| hasher.update((pid as u32).to_be_bytes()));
    sid_i_list.iter().for_each(|sid| hasher.update(sid));

    hasher.finalize().into()
}

pub fn encrypt_message<R: CryptoRngCore>(
    sender_secret_info: (&SecretKey, u8),
    receiver_public_info: (&PublicKey, u8),
    message: &[u8; 64],
    rng: &mut R,
) -> Option<EncryptedData> {
    let sender_box = SalsaBox::new(receiver_public_info.0, sender_secret_info.0);
    let nonce = SalsaBox::generate_nonce(rng);
    sender_box
        .encrypt(&nonce, message.as_slice())
        .ok()
        .and_then(|data| {
            Some(EncryptedData::new(
                data.try_into().ok()?,
                nonce.into(),
                sender_secret_info.1,
                receiver_public_info.1,
            ))
        })
}

pub fn decrypt_message(
    receiver_private_key: &SecretKey,
    sender_public_key: &PublicKey,
    enc_data: &EncryptedData,
) -> Option<[u8; 64]> {
    let receiver_box = SalsaBox::new(sender_public_key, receiver_private_key);
    receiver_box
        .decrypt(
            &GenericArray::from(enc_data.nonce),
            enc_data.ciphertext.as_slice(),
        )
        .ok()
        .and_then(|data| data.try_into().ok())
}

#[cfg(any(test, feature = "test-support"))]
pub mod support {

    use crate::keygen::{utils::setup_keygen, Keyshare};

    use crate::common::{
        ser::Serializable,
        traits::{GroupElem, Round, ScalarReduce},
    };

    /// Execute one round of DKG protocol locally, execute parties in parallel
    /// Used for testing purposes.
    #[allow(clippy::map_identity)]
    pub fn run_round<I, R, O, E>(actors: impl IntoIterator<Item = R>, msgs: I) -> Vec<O>
    where
        R: Round<Input = I, Output = O, Error = E> + Serializable + Send,
        I: Clone + Sync,
        E: std::fmt::Debug,
        O: Send,
    {
        actors
            .into_iter()
            .map(|actor| {
                #[cfg(all(feature = "serde", test))]
                let actor: R = {
                    let mut v = vec![];
                    ciborium::into_writer(&actor, &mut v).unwrap();
                    ciborium::from_reader(v.as_ref() as &[u8]).unwrap()
                };

                actor
            })
            .map(|actor| actor.process(msgs.clone()).unwrap())
            .collect()
    }

    /// Utility function to run the keygen protocol.
    pub fn run_keygen<const T: usize, const N: usize, G: GroupElem>() -> [Keyshare<G>; N]
    where
        G::Scalar: ScalarReduce<[u8; 32]>,
        G::Scalar: Serializable,
    {
        let actors = setup_keygen(T as u8, N as u8);

        let (actors, msgs): (Vec<_>, Vec<_>) = run_round(actors, ()).into_iter().unzip();
        let (actors, msgs): (Vec<_>, Vec<_>) = run_round(actors, msgs).into_iter().unzip();

        run_round(actors, msgs)
            .try_into()
            .map_err(|_| panic!("Failed to convert keyshares"))
            .unwrap()
    }
}
