//! Scheme for deriving deterministic secret keys
//!
//! `DerivableSecret` represents a secret key that can be used to derive child
//! secret keys. A root key secret can be used to derives child
//! keys from it, which can have child keys derived from them, recursively.
//!
//! The `DerivableSecret` struct in this implementation is only used for
//! deriving secret keys, not public keys. This allows supporting multiple
//! crypto schemes for the different cryptographic operations used across the
//! different modules:
//!
//! * secp256k1 for bitcoin deposit addresses, redeem keys and contract keys for
//!   lightning,
//! * bls12-381 for the guardians' threshold signature scheme,
//! * chacha20-poly1305 for symmetric encryption used for backups.
use std::fmt::Formatter;

use bls12_381::Scalar;
use fedimint_core::config::FederationId;
use fedimint_core::encoding::{Decodable, DecodeError, Encodable};
use fedimint_core::module::registry::ModuleDecoderRegistry;
use fedimint_core::secp256k1::{Keypair, Secp256k1, Signing};
use hkdf::hashes::Sha512;
use hkdf::{BitcoinHash, Hkdf};
use ring::aead;

const CHILD_TAG: &[u8; 8] = b"childkey";
const SECP256K1_TAG: &[u8; 8] = b"secp256k";
const BLS12_381_TAG: &[u8; 8] = b"bls12381";
const CHACHA20_POLY1305: &[u8; 8] = b"c20p1305";
const RAW_BYTES: &[u8; 8] = b"rawbytes";

/// Describes a child key of a [`DerivableSecret`]
#[derive(Debug, Copy, Clone, Encodable, Decodable)]
pub struct ChildId(pub u64);

/// A secret that can have child-subkey derived from it.
#[derive(Clone)]
pub struct DerivableSecret {
    /// Derivation level, root = 0, every `child_key` increments it
    level: usize,
    /// An instance of the HKDF (Hash-based Key Derivation
    ///   Function) with SHA-512 as the underlying hash function. It is used to
    ///   derive child keys.
    // TODO: wrap in some secret protecting wrappers maybe?
    kdf: Hkdf<Sha512>,
}

impl DerivableSecret {
    /// Derive root secret key from a secret material and salt.
    ///
    /// The `salt` is just additional data t used
    /// as an additional input to the HKDF.
    pub fn new_root(root_key: &[u8], salt: &[u8]) -> Self {
        DerivableSecret {
            level: 0,
            kdf: Hkdf::new(root_key, Some(salt)),
        }
    }

    /// Get derivation level
    ///
    /// This is useful for ensuring a correct derivation level is used,
    /// in various places.
    ///
    /// Root keys start at `0`, and every derived key increments it.
    pub fn level(&self) -> usize {
        self.level
    }

    pub fn child_key(&self, cid: ChildId) -> DerivableSecret {
        DerivableSecret {
            level: self.level + 1,
            kdf: Hkdf::from_prk(self.kdf.derive_hmac(&tagged_derive(CHILD_TAG, cid))),
        }
    }

    /// Derive a tweaked child key from self using arbitrary bytes.
    ///
    /// Similar to [`child_key`](Self::child_key) but accepts arbitrary bytes
    /// as the tweak input instead of a [`ChildId`]. This increments the
    /// derivation level since it produces a new derived key.
    pub fn tweak(&self, tweak: &[u8]) -> DerivableSecret {
        DerivableSecret {
            level: self.level + 1,
            kdf: Hkdf::from_prk(self.kdf.derive_hmac(tweak)),
        }
    }

    /// Derive a federation-ID-based child key from self.
    ///
    /// This is useful to ensure that the same root secret is not reused
    /// across multiple `fedimint-client` instances for different federations.
    ///
    /// We reset the level to 0 here since `fedimint-client` expects its root
    /// secret to be at that level.
    pub fn federation_key(&self, federation_id: &FederationId) -> DerivableSecret {
        DerivableSecret {
            level: 0,
            kdf: Hkdf::from_prk(
                self.kdf.derive_hmac(&tagged_derive(
                    &federation_id.0.to_byte_array()[..8]
                        .try_into()
                        .expect("Slice with length 8"),
                    ChildId(0),
                )),
            ),
        }
    }

    /// secp256k1 keys are used for bitcoin deposit addresses, redeem keys and
    /// contract keys for lightning.
    pub fn to_secp_key<C: Signing>(self, ctx: &Secp256k1<C>) -> Keypair {
        for key_try in 0u64.. {
            let secret = self
                .kdf
                .derive::<32>(&tagged_derive(SECP256K1_TAG, ChildId(key_try)));
            // The secret not forming a valid key is highly unlikely, this approach is the
            // same used when generating a random secp key.
            if let Ok(key) = Keypair::from_seckey_slice(ctx, &secret) {
                return key;
            }
        }

        unreachable!("If key generation fails this often something else has to be wrong.")
    }

    /// bls12-381 keys are used for the guardians' threshold signature scheme,
    /// and most importantly for its use for the blinding keys for e-cash notes.
    pub fn to_bls12_381_key(&self) -> Scalar {
        Scalar::from_bytes_wide(&self.kdf.derive(&tagged_derive(BLS12_381_TAG, ChildId(0))))
    }

    // `ring` does not support any way to get raw bytes from a key,
    // so we need to be able to get just the raw bytes here, so we can serialize
    // them, and convert to ring type from it.
    pub fn to_chacha20_poly1305_key_raw(&self) -> [u8; 32] {
        self.kdf
            .derive::<32>(&tagged_derive(CHACHA20_POLY1305, ChildId(0)))
    }

    pub fn to_chacha20_poly1305_key(&self) -> aead::UnboundKey {
        aead::UnboundKey::new(
            &aead::CHACHA20_POLY1305,
            &self.to_chacha20_poly1305_key_raw(),
        )
        .expect("created key")
    }

    /// Generate a pseudo-random byte array from the derivable secret.
    pub fn to_random_bytes<const LEN: usize>(&self) -> [u8; LEN] {
        self.kdf.derive(&tagged_derive(RAW_BYTES, ChildId(0)))
    }
}

fn tagged_derive(tag: &[u8; 8], derivation: ChildId) -> [u8; 16] {
    let mut derivation_info = [0u8; 16];
    derivation_info[0..8].copy_from_slice(&tag[..]);
    // The endianness isn't important here because we just need some bytes, but
    // let's use the default for this project (big endian)
    derivation_info[8..16].copy_from_slice(&derivation.0.to_be_bytes()[..]);
    derivation_info
}

impl std::fmt::Debug for DerivableSecret {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "DerivableSecret#")?;
        fedimint_core::format_hex(
            &self
                .kdf
                .derive::<8>(b"just a debug fingerprint derivation salt"),
            f,
        )
    }
}

impl Encodable for DerivableSecret {
    fn consensus_encode<W: std::io::Write>(&self, writer: &mut W) -> Result<(), std::io::Error> {
        (self.level as u64).consensus_encode(writer)?;
        self.kdf.to_prk_bytes().consensus_encode(writer)
    }
}

impl Decodable for DerivableSecret {
    fn consensus_decode_partial<R: std::io::Read>(
        reader: &mut R,
        modules: &ModuleDecoderRegistry,
    ) -> Result<Self, DecodeError> {
        let level_u64 = u64::consensus_decode_partial(reader, modules)?;
        let level = level_u64
            .try_into()
            .map_err(|_| DecodeError::from_str("DerivableSecret level out of range"))?;
        let prk = <[u8; 64]>::consensus_decode_partial(reader, modules)?;

        Ok(Self {
            level,
            kdf: Hkdf::from_prk_bytes(prk),
        })
    }
}

#[cfg(test)]
mod tests {
    use super::{ChildId, Decodable, DerivableSecret, Encodable};

    #[test]
    fn consensus_roundtrip_preserves_derivation() {
        let original = DerivableSecret::new_root(b"root key", b"salt")
            .child_key(ChildId(1))
            .child_key(ChildId(2));

        let encoded = original.consensus_encode_to_vec();
        let decoded = DerivableSecret::consensus_decode_whole(&encoded, &Default::default())
            .expect("decode should succeed");

        assert_eq!(original.level(), decoded.level());
        assert_eq!(
            original.to_random_bytes::<32>(),
            decoded.to_random_bytes::<32>()
        );
        assert_eq!(
            original.child_key(ChildId(3)).to_random_bytes::<32>(),
            decoded.child_key(ChildId(3)).to_random_bytes::<32>()
        );
    }
}
