// Copyright 2019 The Tari Project
//
// Redistribution and use in source and binary forms, with or without modification, are permitted provided that the
// following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following
// disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the
// following disclaimer in the documentation and/or other materials provided with the distribution.
//
// 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote
// products derived from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
// INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
// USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

use std::marker::PhantomData;

use derivative::Derivative;
use digest::{Digest, consts::U64, typenum::IsEqual};
use serde::{Deserialize, Serialize};
use tari_common_types::{
    seeds::cipher_seed::CipherSeed,
    types::{CompressedPublicKey, PrivateKey},
};
use tari_crypto::{
    hashing::{DomainSeparatedHasher, LengthExtensionAttackResistant},
    keys::SecretKey,
    tari_utilities::byte_array::ByteArrayError,
};
use tari_hashing::KeyManagerDomain;
use zeroize::Zeroize;

use crate::legacy_key_manager::HASHER_LABEL_DERIVE_KEY;

#[derive(Clone, Derivative, Serialize, Deserialize, Zeroize)]
#[derivative(Debug)]
pub struct DerivedKey {
    #[derivative(Debug = "ignore")]
    #[serde(skip_deserializing)]
    pub key: PrivateKey,
    pub key_index: u64,
}

#[derive(Clone, Derivative, Serialize, Deserialize, Zeroize)]
#[derivative(Debug)]
pub struct DerivedPublicKey {
    #[derivative(Debug = "ignore")]
    #[serde(skip_deserializing)]
    pub key: CompressedPublicKey,
    pub key_index: u64,
}

#[derive(Clone, Derivative, PartialEq, Serialize, Deserialize, Zeroize)]
#[derivative(Debug)]
pub struct TariKeyManager<D: Digest + LengthExtensionAttackResistant> {
    #[derivative(Debug = "ignore")]
    seed: CipherSeed,
    #[derivative(Debug = "ignore")]
    pub branch_seed: String,
    primary_key_index: u64,
    digest_type: PhantomData<D>,
}

impl<D> TariKeyManager<D>
where
    D: Digest + LengthExtensionAttackResistant,
    D::OutputSize: IsEqual<U64>,
{
    /// Creates a new KeyManager with a new randomly selected entropy
    pub fn random() -> TariKeyManager<D> {
        TariKeyManager {
            seed: CipherSeed::random(),
            branch_seed: "".to_string(),
            primary_key_index: 0,
            digest_type: PhantomData,
        }
    }

    /// Constructs a KeyManager from known parts
    pub fn from(seed: CipherSeed, branch_seed: String, primary_key_index: u64) -> TariKeyManager<D> {
        TariKeyManager {
            seed,
            branch_seed,
            primary_key_index,
            digest_type: PhantomData,
        }
    }

    /// Derive a new private key from master key: derived_key=H(master_key||branch_seed||index), for some
    /// hash function H which is Length attack resistant, such as Blake2b.
    fn derive_private_key(&self, key_index: u64) -> Result<PrivateKey, ByteArrayError> {
        // apply domain separation to generate derive key. Under the hood, the hashing api prepends the length of each
        // piece of data for concatenation, reducing the risk of collisions due to redundancy of variable length
        // input
        let derive_key = DomainSeparatedHasher::<D, KeyManagerDomain>::new_with_label(HASHER_LABEL_DERIVE_KEY)
            .chain(self.seed.entropy())
            .chain(self.branch_seed.as_bytes())
            .chain(key_index.to_le_bytes())
            .finalize();

        let derive_key = derive_key.as_ref();
        let s = PrivateKey::from_uniform_bytes(derive_key)?;
        Ok(s)
    }

    /// Derive a new private key from master key: derived_key=H(master_key||branch_seed||index), for some
    /// hash function H which is Length attack resistant, such as Blake2b.
    pub fn derive_key(&self, key_index: u64) -> Result<DerivedKey, ByteArrayError> {
        let secret = self.derive_private_key(key_index)?;
        Ok(DerivedKey { key: secret, key_index })
    }

    /// Derive a new public key from master key: derived_key=H(master_key||branch_seed||index), for some
    /// hash function H which is Length attack resistant, such as Blake2b.
    pub fn derive_public_key(&self, key_index: u64) -> Result<DerivedPublicKey, ByteArrayError> {
        let secret = self.derive_private_key(key_index)?;
        Ok(DerivedPublicKey {
            key: CompressedPublicKey::from_secret_key(&secret),
            key_index,
        })
    }

    pub fn get_private_key(&self, key_index: u64) -> Result<PrivateKey, ByteArrayError> {
        let secret = self.derive_private_key(key_index)?;
        Ok(secret)
    }

    /// Generate next deterministic private key derived from master key
    pub fn next_key(&mut self) -> Result<DerivedKey, ByteArrayError> {
        self.primary_key_index += 1;
        self.derive_key(self.primary_key_index)
    }

    /// Generate next deterministic private key derived from master key
    pub fn increment_key_index(&mut self, increment: u64) -> u64 {
        self.primary_key_index += increment;
        self.primary_key_index
    }

    pub fn key_index(&self) -> u64 {
        self.primary_key_index
    }

    pub fn update_key_index(&mut self, new_index: u64) {
        self.primary_key_index = new_index;
    }
}

#[cfg(test)]
mod test {
    use std::str::FromStr;

    use blake2::Blake2b;
    use tari_transaction_components::key_manager::{SecretTransactionKeyManagerInterface, TariKeyId};

    use super::*;
    use crate::legacy_key_manager::{
        LegacySerializedKeyString,
        LegacyTariKeyId,
        LegacyTransactionKeyManagerInterface,
        create_new_random_key_manager,
    };

    #[test]
    fn test_new_keymanager() {
        let km1 = TariKeyManager::<Blake2b<U64>>::random();
        let km2 = TariKeyManager::<Blake2b<U64>>::random();
        assert_ne!(km1.seed, km2.seed);
    }

    #[test]
    fn test_derive_and_next_key() {
        let mut km = TariKeyManager::<Blake2b<U64>>::random();
        let next_key1_result = km.next_key();
        let next_key2_result = km.next_key();
        let desired_key_index1 = 1;
        let desired_key_index2 = 2;
        let derived_key1_result = km.derive_key(desired_key_index1);
        let derived_key2_result = km.derive_key(desired_key_index2);
        let next_key1 = next_key1_result.unwrap();
        let next_key2 = next_key2_result.unwrap();
        let derived_key1 = derived_key1_result.unwrap();
        let derived_key2 = derived_key2_result.unwrap();
        assert_ne!(next_key1.key, next_key2.key);
        assert_eq!(next_key1.key, derived_key1.key);
        assert_eq!(next_key2.key, derived_key2.key);
        assert_eq!(next_key1.key_index, desired_key_index1);
        assert_eq!(next_key2.key_index, desired_key_index2);
    }

    #[test]
    fn test_derive_and_next_key_with_branch_seed() {
        let mut km = TariKeyManager::<Blake2b<U64>>::from(CipherSeed::random(), "Test".to_string(), 0);
        let next_key1_result = km.next_key();
        let next_key2_result = km.next_key();
        let desired_key_index1 = 1;
        let desired_key_index2 = 2;
        let derived_key1_result = km.derive_key(desired_key_index1);
        let derived_key2_result = km.derive_key(desired_key_index2);
        let next_key1 = next_key1_result.unwrap();
        let next_key2 = next_key2_result.unwrap();
        let derived_key1 = derived_key1_result.unwrap();
        let derived_key2 = derived_key2_result.unwrap();
        assert_ne!(next_key1.key, next_key2.key);
        assert_eq!(next_key1.key, derived_key1.key);
        assert_eq!(next_key2.key, derived_key2.key);
        assert_eq!(next_key1.key_index, desired_key_index1);
        assert_eq!(next_key2.key_index, desired_key_index2);
    }

    #[test]
    fn test_use_of_branch_seed() {
        let x = CipherSeed::random();
        let mut km1 = TariKeyManager::<Blake2b<U64>>::from(x.clone(), "some".to_string(), 0);
        let mut km2 = TariKeyManager::<Blake2b<U64>>::from(x, "other".to_string(), 0);
        let next_key1 = km1.next_key().unwrap();
        let next_key2 = km2.next_key().unwrap();
        assert_ne!(next_key1.key, next_key2.key);
    }

    #[tokio::test]
    async fn test_convert_legacy_managed_and_derived_key() {
        let legacy_key_manager = create_new_random_key_manager().await.unwrap();

        // Managed key
        let legacy_managed = LegacyTariKeyId::Managed {
            branch: "commitment mask".to_string(),
            index: 18309514604232961632,
        };
        // Derived key wrapping managed
        let legacy_derived = LegacyTariKeyId::Derived {
            key: LegacySerializedKeyString::from(legacy_managed.to_string()),
        };

        let legacy_managed_converted = legacy_key_manager
            .convert_legacy_tari_key_id_to_current(&legacy_managed)
            .unwrap();
        let legacy_derived_converted = legacy_key_manager
            .convert_legacy_tari_key_id_to_current(&legacy_derived)
            .unwrap();

        let legacy_managed_private_key = legacy_key_manager.get_private_key(&legacy_managed_converted).unwrap();
        let legacy_derived_private_key = legacy_key_manager.get_private_key(&legacy_derived_converted).unwrap();
        //  Legacy managed and derived private keys should be different
        assert_ne!(legacy_managed_private_key, legacy_derived_private_key);

        let key_manager = legacy_key_manager.key_manager();

        let managed_private_key = key_manager.get_private_key(&legacy_managed_converted).unwrap();
        let derived_private_key = key_manager.get_private_key(&legacy_derived_converted).unwrap();
        // New managed and derived private keys should be different
        assert_ne!(managed_private_key, derived_private_key);

        // Legacy and new managed_private_key should be the same
        assert_eq!(legacy_managed_private_key, managed_private_key);
        // Legacy and new derived_private_key should be the same
        assert_eq!(legacy_derived_private_key, derived_private_key);

        // These must fail because the legacy format is not compatible with the current TariKeyId parser
        assert!(TariKeyId::from_str(&legacy_managed.to_string()).is_err());
        assert!(TariKeyId::from_str(&legacy_derived.to_string()).is_err());

        // These must pass because they have been converted to the current format
        assert!(TariKeyId::from_str(&legacy_managed_converted.to_string()).is_ok());
        assert!(TariKeyId::from_str(&legacy_derived_converted.to_string()).is_ok());
    }
}
