use crate::{DerivationIndex, PublicKey};

use hmac::{Hmac, Mac};
use sha2::Sha512;
use thiserror::Error;

const KEY_SIZE: usize = 256;

/// Errors associated with creating or deriving Extended Public Keys.
#[derive(Error, Clone, Debug)]
pub enum Error {
    /// BIP-0032 does not support deriving public keys from hardened private keys.
    #[error("Public keys may not be derived when hardened")]
    HardenedIndex,

    /// The maximum depth for key derivation has been reached. The supported maximum depth is 255.
    #[error("Reached maximum depth for key derivation")]
    Depth,

    /// This should never happen. HMAC keys should be able to take any size input.
    #[error("Invalid length for HMAC key while generating master key (report me!)")]
    HmacInvalidLength(#[from] hmac::digest::InvalidLength),

    /// An unknown error occurred while deriving a child key.
    #[error("Unknown error while deriving child key")]
    Derivation,
}

pub type Result<T, E = Error> = std::result::Result<T, E>;
pub type ChainCode = [u8; 32];
type HmacSha512 = Hmac<Sha512>;

/// Extended public keys derived using BIP-0032.
///
/// Generic over types implementing [`PublicKey`].
pub struct ExtendedPublicKey<K: PublicKey> {
    public_key: K,
    depth: u8,
    chain_code: ChainCode,
}

impl<K> ExtendedPublicKey<K>
where
    K: PublicKey,
{
    /*
    /// Create a new [`ExtendedPublicKey`] from previously known values.
    pub fn new(public_key: K, chain_code: ChainCode) -> Self {
        Self {
            public_key,
            depth: 0,
            chain_code,
        }
    }
    */

    /// Derive a child with a given [`DerivationIndex`].
    ///
    /// # Panics
    ///
    /// The method performs unchecked `try_into()` operations on a constant-sized slice.
    ///
    /// # Errors
    ///
    /// An error may be returned if:
    ///
    /// * The depth exceeds the maximum depth [`u8::MAX`].
    /// * A `HmacSha512` can't be constructed - this should be impossible.
    /// * Deriving a child key fails. Check the documentation for your [`PublicKey`].
    pub fn derive_child(&self, index: &DerivationIndex) -> Result<Self> {
        if index.is_hardened() {
            return Err(Error::HardenedIndex);
        }

        let depth = self.depth.checked_add(1).ok_or(Error::Depth)?;

        let hmac = HmacSha512::new_from_slice(&self.chain_code)
            .map_err(Error::from)?
            .chain_update(self.public_key.to_bytes())
            .chain_update(index.to_bytes())
            .finalize()
            .into_bytes();

        let (child_key, chain_code) = hmac.split_at(KEY_SIZE / 8);
        let derived_key = self
            .public_key
            .derive_child(child_key.try_into().expect("Invalid key length"))
            .map_err(|_| Error::Derivation)?;
        let chain_code = chain_code.try_into().expect("Invalid chain code length");

        Ok(Self {
            public_key: derived_key,
            depth,
            chain_code,
        })
    }
}