/*
 *
 *    Copyright (c) 2020-2022 Project CHIP Authors
 *
 *    Licensed under the Apache License, Version 2.0 (the "License");
 *    you may not use this file except in compliance with the License.
 *    You may obtain a copy of the License at
 *
 *        http://www.apache.org/licenses/LICENSE-2.0
 *
 *    Unless required by applicable law or agreed to in writing, software
 *    distributed under the License is distributed on an "AS IS" BASIS,
 *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *    See the License for the specific language governing permissions and
 *    limitations under the License.
 */

use std::convert::{TryFrom, TryInto};

use aes::Aes128;
use ccm::{
    aead::generic_array::GenericArray,
    consts::{U13, U16},
    Ccm,
};
use elliptic_curve::sec1::{FromEncodedPoint, ToEncodedPoint};
use hmac::Mac;
use log::error;
use p256::{
    ecdsa::{Signature, SigningKey, VerifyingKey},
    AffinePoint, EncodedPoint, PublicKey, SecretKey,
};
use sha2::Digest;
use x509_cert::{
    attr::AttributeType,
    der::{asn1::BitString, Any, Encode},
    name::RdnSequence,
    request::CertReq,
    spki::{AlgorithmIdentifier, SubjectPublicKeyInfoOwned},
};

use crate::error::Error;

use super::CryptoKeyPair;

type HmacSha256I = hmac::Hmac<sha2::Sha256>;
type AesCcm = Ccm<Aes128, U16, U13>;

#[derive(Clone)]
pub struct Sha256 {
    hasher: sha2::Sha256,
}

impl Sha256 {
    pub fn new() -> Result<Self, Error> {
        Ok(Self {
            hasher: sha2::Sha256::new(),
        })
    }

    pub fn update(&mut self, data: &[u8]) -> Result<(), Error> {
        self.hasher.update(data);
        Ok(())
    }

    pub fn finish(self, digest: &mut [u8]) -> Result<(), Error> {
        let output = self.hasher.finalize();
        digest.copy_from_slice(output.as_slice());
        Ok(())
    }
}

pub struct HmacSha256 {
    inner: HmacSha256I,
}

impl HmacSha256 {
    pub fn new(key: &[u8]) -> Result<Self, Error> {
        Ok(Self {
            inner: HmacSha256I::new_from_slice(key).map_err(|e| {
                error!("Error creating HmacSha256 {:?}", e);
                Error::TLSStack
            })?,
        })
    }

    pub fn update(&mut self, data: &[u8]) -> Result<(), Error> {
        self.inner.update(data);
        Ok(())
    }

    pub fn finish(self, out: &mut [u8]) -> Result<(), Error> {
        let result = &self.inner.finalize().into_bytes()[..];
        out.clone_from_slice(result);
        Ok(())
    }
}

pub enum KeyType {
    Private(SecretKey),
    Public(PublicKey),
}

pub struct KeyPair {
    key: KeyType,
}

impl KeyPair {
    pub fn new() -> Result<Self, Error> {
        let mut rng = rand::thread_rng();
        let secret_key = SecretKey::random(&mut rng);

        Ok(Self {
            key: KeyType::Private(secret_key),
        })
    }

    pub fn new_from_components(pub_key: &[u8], priv_key: &[u8]) -> Result<Self, Error> {
        let secret_key = SecretKey::from_slice(priv_key).unwrap();
        let encoded_point = EncodedPoint::from_bytes(pub_key).unwrap();
        let public_key = PublicKey::from_encoded_point(&encoded_point).unwrap();
        assert_eq!(public_key, secret_key.public_key());

        Ok(Self {
            key: KeyType::Private(secret_key),
        })
    }

    pub fn new_from_public(pub_key: &[u8]) -> Result<Self, Error> {
        let encoded_point = EncodedPoint::from_bytes(pub_key).unwrap();
        Ok(Self {
            key: KeyType::Public(PublicKey::from_encoded_point(&encoded_point).unwrap()),
        })
    }

    fn public_key_point(&self) -> AffinePoint {
        match &self.key {
            KeyType::Private(k) => *(k.public_key().as_affine()),
            KeyType::Public(k) => *(k.as_affine()),
        }
    }

    fn private_key(&self) -> Result<&SecretKey, Error> {
        match &self.key {
            KeyType::Private(key) => Ok(key),
            KeyType::Public(_) => Err(Error::Crypto),
        }
    }
}

impl CryptoKeyPair for KeyPair {
    fn get_private_key(&self, priv_key: &mut [u8]) -> Result<usize, Error> {
        match &self.key {
            KeyType::Private(key) => {
                let bytes = key.to_bytes();
                let slice = bytes.as_slice();
                let len = slice.len();
                priv_key[..slice.len()].copy_from_slice(slice);
                Ok(len)
            }
            KeyType::Public(_) => Err(Error::Crypto),
        }
    }
    fn get_csr<'a>(&self, out_csr: &'a mut [u8]) -> Result<&'a [u8], Error> {
        use p256::ecdsa::signature::Signer;

        let subject = RdnSequence(vec![x509_cert::name::RelativeDistinguishedName(
            vec![x509_cert::attr::AttributeTypeAndValue {
                // Organization name: http://www.oid-info.com/get/2.5.4.10
                oid: x509_cert::attr::AttributeType::new_unwrap("2.5.4.10"),
                value: x509_cert::attr::AttributeValue::new(
                    x509_cert::der::Tag::Utf8String,
                    "CSR".as_bytes(),
                )
                .unwrap(),
            }]
            .try_into()
            .unwrap(),
        )]);
        let mut pubkey = [0; 65];
        self.get_public_key(&mut pubkey).unwrap();
        let info = x509_cert::request::CertReqInfo {
            version: x509_cert::request::Version::V1,
            subject,
            public_key: SubjectPublicKeyInfoOwned {
                algorithm: AlgorithmIdentifier {
                    // ecPublicKey(1) http://www.oid-info.com/get/1.2.840.10045.2.1
                    oid: AttributeType::new_unwrap("1.2.840.10045.2.1"),
                    parameters: Some(
                        Any::new(
                            x509_cert::der::Tag::ObjectIdentifier,
                            // prime256v1 http://www.oid-info.com/get/1.2.840.10045.3.1.7
                            AttributeType::new_unwrap("1.2.840.10045.3.1.7").as_bytes(),
                        )
                        .unwrap(),
                    ),
                },
                subject_public_key: BitString::from_bytes(&pubkey).unwrap(),
            },
            attributes: Default::default(),
        };
        let mut message = vec![];
        info.encode(&mut message).unwrap();

        // Can't use self.sign_msg as the signature has to be in DER format
        let private_key = self.private_key()?;
        let signing_key = SigningKey::from(private_key);
        let sig: Signature = signing_key.sign(&message);
        let to_der = sig.to_der();
        let signature = to_der.as_bytes();

        let cert = CertReq {
            info,
            algorithm: AlgorithmIdentifier {
                // ecdsa-with-SHA256(2) http://www.oid-info.com/get/1.2.840.10045.4.3.2
                oid: AttributeType::new_unwrap("1.2.840.10045.4.3.2"),
                parameters: None,
            },
            signature: BitString::from_bytes(signature).unwrap(),
        };
        let out = cert.to_der().unwrap();
        let a = &mut out_csr[0..out.len()];
        a.copy_from_slice(&out);

        Ok(a)
    }
    fn get_public_key(&self, pub_key: &mut [u8]) -> Result<usize, Error> {
        let point = self.public_key_point().to_encoded_point(false);
        let bytes = point.as_bytes();
        let len = bytes.len();
        pub_key[..len].copy_from_slice(bytes);
        Ok(len)
    }
    fn derive_secret(self, peer_pub_key: &[u8], secret: &mut [u8]) -> Result<usize, Error> {
        let encoded_point = EncodedPoint::from_bytes(peer_pub_key).unwrap();
        let peer_pubkey = PublicKey::from_encoded_point(&encoded_point).unwrap();
        let private_key = self.private_key()?;
        let shared_secret = elliptic_curve::ecdh::diffie_hellman(
            private_key.to_nonzero_scalar(),
            peer_pubkey.as_affine(),
        );
        let bytes = shared_secret.raw_secret_bytes();
        let bytes = bytes.as_slice();
        let len = bytes.len();
        assert_eq!(secret.len(), len);
        secret.copy_from_slice(bytes);

        Ok(len)
    }
    fn sign_msg(&self, msg: &[u8], signature: &mut [u8]) -> Result<usize, Error> {
        use p256::ecdsa::signature::Signer;

        if signature.len() < super::EC_SIGNATURE_LEN_BYTES {
            return Err(Error::NoSpace);
        }

        match &self.key {
            KeyType::Private(k) => {
                let signing_key = SigningKey::from(k);
                let sig: Signature = signing_key.sign(msg);
                let bytes = sig.to_bytes().to_vec();
                let len = bytes.len();
                signature[..len].copy_from_slice(&bytes);
                Ok(len)
            }
            KeyType::Public(_) => todo!(),
        }
    }
    fn verify_msg(&self, msg: &[u8], signature: &[u8]) -> Result<(), Error> {
        use p256::ecdsa::signature::Verifier;

        let verifying_key = VerifyingKey::from_affine(self.public_key_point()).unwrap();
        let signature = Signature::try_from(signature).unwrap();

        verifying_key
            .verify(msg, &signature)
            .map_err(|_| Error::InvalidSignature)?;

        Ok(())
    }
}

pub fn pbkdf2_hmac(pass: &[u8], iter: usize, salt: &[u8], key: &mut [u8]) -> Result<(), Error> {
    pbkdf2::pbkdf2::<hmac::Hmac<sha2::Sha256>>(pass, salt, iter as u32, key).unwrap();

    Ok(())
}

pub fn hkdf_sha256(salt: &[u8], ikm: &[u8], info: &[u8], key: &mut [u8]) -> Result<(), Error> {
    hkdf::Hkdf::<sha2::Sha256>::new(Some(salt), ikm)
        .expand(info, key)
        .map_err(|e| {
            error!("Error with hkdf_sha256 {:?}", e);
            Error::TLSStack
        })
}

// TODO: add tests and check against mbedtls and openssl
pub fn encrypt_in_place(
    key: &[u8],
    nonce: &[u8],
    ad: &[u8],
    data: &mut [u8],
    data_len: usize,
) -> Result<usize, Error> {
    use ccm::{AeadInPlace, KeyInit};

    let key = GenericArray::from_slice(key);
    let nonce = GenericArray::from_slice(nonce);
    let cipher = AesCcm::new(key);
    // This is probably incorrect
    let mut buffer = data[0..data_len].to_vec();
    cipher.encrypt_in_place(nonce, ad, &mut buffer)?;
    let len = buffer.len();
    data.clone_from_slice(&buffer[..]);

    Ok(len)
}

pub fn decrypt_in_place(
    key: &[u8],
    nonce: &[u8],
    ad: &[u8],
    data: &mut [u8],
) -> Result<usize, Error> {
    use ccm::{AeadInPlace, KeyInit};

    let key = GenericArray::from_slice(key);
    let nonce = GenericArray::from_slice(nonce);
    let cipher = AesCcm::new(key);
    // This is probably incorrect
    let mut buffer = data.to_vec();
    cipher.decrypt_in_place(nonce, ad, &mut buffer)?;
    let len = buffer.len();
    data[..len].copy_from_slice(&buffer[..]);

    Ok(len)
}