use crate::*;
#[encrypted_library]
mod arcis_library {

    /// Owner for shared data between the owner of a public key and the MXE.
    ///
    /// Decrypting data owned by this owner this requires participation from either
    /// * the whole MXE
    /// * the public key owner.
    #[derive(Debug, PartialEq)]
    pub struct Shared {
        pub public_key: ArcisX25519Pubkey,
        pub(crate) nonce: u128,
    }
    impl Shared {
        /// Creates an owner for shared data between the MXE and the public key owner.
        pub fn new(public_key: ArcisX25519Pubkey) -> Self {
            let nonce = ArcisRNG::gen_public_integer_from_width(128);
            Self { public_key, nonce }
        }
    }
    /// Owner of secret data.
    /// Decrypting data owned by this owner requires participation from all the nodes that comprises
    /// the cluster attached to the given MXE.
    #[derive(Debug, PartialEq)]
    pub struct Mxe {
        pub(crate) nonce: u128,
    }
    impl Mxe {
        /// Generate a nonce for the MXE. Used to encrypt data which can only be decrypted by the
        /// nodes that comprise the cluster attached to the given MXE.
        pub fn get() -> Self {
            let nonce = ArcisRNG::gen_public_integer_from_width(128);
            Self { nonce }
        }
    }

    pub(crate) fn base58_to_32_uint8array(s: &[u8]) -> [u8; 32] {
        fn base58_to_u8(c: u8) -> u8 {
            #![allow(clippy::manual_range_contains)]
            if b'1' <= c && c <= b'9' {
                c - b'1'
            } else if b'A' <= c && c <= b'H' {
                c - b'A' + 9
            } else if b'J' <= c && c <= b'N' {
                c - b'J' + 17
            } else if b'P' <= c && c <= b'Z' {
                c - b'P' + 22
            } else if b'a' <= c && c <= b'k' {
                c - b'a' + 33
            } else if b'm' <= c && c <= b'z' {
                c - b'm' + 44
            } else {
                // Maybe some kind of assert or panic would be good.
                127
            }
        }
        fn op(a: &mut [u8; 32], factor: u8, carry: u8) {
            let mut carry = carry as u16;
            for a in a.iter_mut() {
                let num = *a as u16 * factor as u16 + carry;
                carry = num / 256;
                *a = (num % 256) as u8;
            }
        }
        let mut address = [0u8; 32];
        for &c in s {
            let b = base58_to_u8(c);
            op(&mut address, 58, b);
        }
        address.reverse();
        address
    }

    impl ArcisX25519Pubkey {
        /// Creates an Arcis public key based on the base58-encoding of the Montgomery x-coordinate.
        pub fn from_base58(s: &[u8]) -> Self {
            let arr = base58_to_32_uint8array(s);
            Self::from_uint8(&arr)
        }
        /// Creates an Arcis public key based on the uint8array-encoding of the Montgomery
        /// x-coordinate.
        pub fn from_uint8(s: &[u8]) -> Self {
            if s.len() > 32 {
                arcis_static_panic!("Length of arg in `from_uint8` is too big.");
            }
            let two_power_eight = BaseField::from(256);
            let mut x = BaseField::from(0);
            let mut factor = BaseField::from(1);
            for &b in s {
                x += BaseField::from(b as u64) * factor;
                factor *= two_power_eight;
            }
            Self::new_from_x(x)
        }
    }

    impl<T: ArcisType> EncData<T> {
        /// This does the same thing as `.to_arcis()` on the matching `Enc`.
        ///
        /// This function exists for performance reasons: imagine you have an `#[instruction]`
        /// with inputs `t: Enc<Shared, T>` and `u: Enc<Shared, U>` that have the same pubkey.
        /// Arcis cannot know the pubkeys are equal and will have duplicate gates for each pubkey.
        /// You have two ways of fixing the performance hit:
        /// * Refactoring your `#[instruction]` inputs into `tu: Enc<Shared, (T, U)>`. But that may
        ///   require to change the smart contract and the whole way your app flows.
        /// * Refactoring your `#[instruction]` inputs into `key: ArcisX25519Pubkey, t_nonce: u128,
        ///   t: EncData<T>, u_nonce: u128, u: EncData<U>`. This change will only require to change
        ///   the caller and the callee.
        /// The first solution is slightly better in terms of performance if you can do it.
        pub fn to_arcis_with_pubkey_and_nonce(
            self,
            public_key: ArcisX25519Pubkey,
            nonce: u128,
        ) -> T {
            let owner = Shared { public_key, nonce };
            Enc { owner, data: self }.to_arcis()
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use arcis_compiler::{
        traits::FromLeBytes,
        utils::{
            crypto::key::{X25519PrivateKey, X25519PublicKey},
            field::ScalarField,
        },
    };

    #[test]
    fn base_58_conversion() {
        // Python output of:
        // keystring_base58="2uKu51kQaLseu7FySMAGWU6hpnjNvgGr3PkvUCBVTTPD"
        // hex_bytes = base58.b58decode(keystring_base58)
        // print(hex_bytes)
        assert_eq!(
            base58_to_32_uint8array(b"2uKu51kQaLseu7FySMAGWU6hpnjNvgGr3PkvUCBVTTPD").as_slice(),
            b"\x1cCA\xba\x1e\xc36$s\xab3m%Q\x0b\x9c\x18\x7f $-z\x9cS\x03M\xb4\x1c\x86\x11T\x00"
        );
        assert_eq!(
            // x25519 pubkey corresponding to the below private key
            ArcisX25519Pubkey::from_uint8(&[
                205, 104, 97, 219, 73, 89, 119, 42, 237, 127, 47, 222, 77, 203, 82, 49, 97, 21,
                242, 44, 104, 77, 109, 141, 78, 77, 25, 54, 179, 176, 75, 13,
            ]),
            ArcisX25519Pubkey::new(
                X25519PublicKey::new_from_private_key(
                    X25519PrivateKey::<ScalarField>::from_le_bytes([
                        159, 131, 199, 155, 47, 64, 100, 24, 186, 230, 173, 56, 87, 137, 173, 118,
                        68, 37, 88, 20, 117, 108, 87, 11, 167, 239, 133, 57, 235, 122, 16, 238,
                    ])
                )
                .inner()
            )
        );
    }
}
