mod bit_accumulator;
mod entropy;
mod phrase;

#[cfg(not(feature = "std"))]
use alloc::{
    borrow::Cow,
    format,
    string::{String, ToString},
    vec::Vec,
};
use core::{fmt, marker::PhantomData, mem, str};
#[cfg(feature = "std")]
use std::borrow::Cow;

use hmac::Hmac;
use sha2::Sha512;
use unicode_normalization::{IsNormalized, UnicodeNormalization, is_nfkd_quick};
use zeroize::Zeroizing;

use self::{
    bit_accumulator::BitAccumulator,
    entropy::encode_entropy,
    phrase::{DecodeMode, decode_phrase},
};
use crate::{
    error::Error,
    language::{English, Language},
};

const BITS_PER_WORD: usize = 11;
const BITS_PER_BYTE: usize = 8;

/// Determines the words count that will be present in a [`Mnemonic`](crate::Mnemonic) phrase.
#[derive(Copy, Clone, Debug, Default, Ord, PartialOrd, Eq, PartialEq, Hash)]
pub enum Count {
    /// 12 words, entropy length: 128 bits, the checksum length: 4 bits.
    #[default]
    Words12,
    /// 15 words, entropy length: 160 bits, the checksum length: 5 bits.
    Words15,
    /// 18 words, entropy length: 192 bits, the checksum length: 6 bits.
    Words18,
    /// 21 words, entropy length: 224 bits, the checksum length: 7 bits.
    Words21,
    /// 24 words, entropy length: 256 bits, the checksum length: 8 bits.
    Words24,
}

impl fmt::Display for Count {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{} words (entropy {} bits + checksum {} bits)",
            self.word_count(),
            self.entropy_bit_length(),
            self.checksum_bit_length()
        )
    }
}

impl From<Count> for usize {
    fn from(count: Count) -> Self {
        match count {
            Count::Words12 => 12,
            Count::Words15 => 15,
            Count::Words18 => 18,
            Count::Words21 => 21,
            Count::Words24 => 24,
        }
    }
}

impl TryFrom<usize> for Count {
    type Error = Error;

    fn try_from(count: usize) -> Result<Self, Self::Error> {
        Self::from_word_count(count)
    }
}

impl Count {
    /// Creates a [`Count`] for a mnemonic phrase with the given word count.
    const fn from_word_count(count: usize) -> Result<Self, Error> {
        Ok(match count {
            12 => Self::Words12,
            15 => Self::Words15,
            18 => Self::Words18,
            21 => Self::Words21,
            24 => Self::Words24,
            others => return Err(Error::BadWordCount(others)),
        })
    }

    /// Creates a [`Count`] for a mnemonic phrase with the given entropy bits size.
    const fn from_key_size(size: usize) -> Result<Self, Error> {
        Ok(match size {
            128 => Self::Words12,
            160 => Self::Words15,
            192 => Self::Words18,
            224 => Self::Words21,
            256 => Self::Words24,
            others => return Err(Error::BadEntropyBitCount(others)),
        })
    }

    /// Creates a [`Count`] for an existing mnemonic phrase.
    fn from_phrase<P: AsRef<str>>(phrase: P) -> Result<Self, Error> {
        let word_count = phrase.as_ref().split_whitespace().count();
        Self::from_word_count(word_count)
    }

    /// Returns the number of words.
    pub const fn word_count(&self) -> usize {
        match self {
            Self::Words12 => 12,
            Self::Words15 => 15,
            Self::Words18 => 18,
            Self::Words21 => 21,
            Self::Words24 => 24,
        }
    }

    /// Returns the size of `entropy + checksum` in bits.
    pub const fn total_bit_length(&self) -> usize {
        self.entropy_bit_length() + self.checksum_bit_length()
    }

    /// Returns the size of entropy in bits.
    pub const fn entropy_bit_length(&self) -> usize {
        match self {
            Self::Words12 => 128,
            Self::Words15 => 160,
            Self::Words18 => 192,
            Self::Words21 => 224,
            Self::Words24 => 256,
        }
    }

    /// Returns the size of checksum in bits.
    pub const fn checksum_bit_length(&self) -> usize {
        match self {
            Self::Words12 => 4,
            Self::Words15 => 5,
            Self::Words18 => 6,
            Self::Words21 => 7,
            Self::Words24 => 8,
        }
    }
}

/// A mnemonic representation.
///
/// First, an initial entropy of ENT bits is generated.
/// A checksum is generated by taking the first `ENT/32` bits of its SHA256 hash.
/// This checksum is appended to the end of the initial entropy.
///
/// Next, these concatenated bits are split into groups of `11` bits,
/// each encoding a number from 0-2047, serving as an index into a wordlist.
///
/// Finally, we convert these numbers into words and use the joined words as a mnemonic sentence.
///
/// - **ENT**: the initial entropy length
/// - **CS**: the checksum length
/// - **MS**: the length of the generated mnemonic sentence in words
///
/// **CS** = **ENT** / 32
///
/// **MS** = (**ENT** + **CS**) / 11
///
/// |  ENT  |  CS  | ENT+CS |  MS  |
/// | :---: | :--: | :----: | :--: |
/// |  128  |  4   |  132   |  12  |
/// |  160  |  5   |  165   |  15  |
/// |  192  |  6   |  198   |  18  |
/// |  224  |  7   |  231   |  21  |
/// |  256  |  8   |  264   |  24  |
///
/// For example, a 12 word mnemonic phrase is essentially a friendly representation of
/// a 128-bit key, while a 24 word mnemonic phrase is essentially a 256-bit key.
#[derive(Clone, Eq, PartialEq)]
pub struct Mnemonic<L = English> {
    lang: PhantomData<L>,
    phrase: Zeroizing<String>,
    entropy: Zeroizing<Vec<u8>>,
}

impl<L: Language> fmt::Debug for Mnemonic<L> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.phrase())
    }
}

impl<L: Language> fmt::Display for Mnemonic<L> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.phrase())
    }
}

impl<L: Language> str::FromStr for Mnemonic<L> {
    type Err = Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        Self::from_phrase(s)
    }
}

impl<L: Language> AsRef<str> for Mnemonic<L> {
    fn as_ref(&self) -> &str {
        self.phrase()
    }
}

impl<L: Language> Mnemonic<L> {
    /// Generates a new [`Mnemonic`] randomly in the specified word count.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use bip0039::{Count, Mnemonic};
    ///
    /// let mnemonic = <Mnemonic>::generate(Count::Words12);
    /// let phrase = mnemonic.phrase();
    /// ```
    #[cfg_attr(
        feature = "chinese-simplified",
        doc = r##"
```rust
use bip0039::{ChineseSimplified, Count, Mnemonic};

let mnemonic = <Mnemonic<ChineseSimplified>>::generate(Count::Words24);
let phrase = mnemonic.phrase();
```
"##
    )]
    #[cfg(feature = "rand")]
    pub fn generate(word_count: Count) -> Self {
        use rand::RngCore;
        let mut rng = rand::rng();

        match word_count {
            Count::Words12 => {
                let mut entropy = [0u8; 16];
                rng.fill_bytes(&mut entropy);
                Self::from_entropy(entropy)
            },
            Count::Words15 => {
                let mut entropy = [0u8; 20];
                rng.fill_bytes(&mut entropy);
                Self::from_entropy(entropy)
            },
            Count::Words18 => {
                let mut entropy = [0u8; 24];
                rng.fill_bytes(&mut entropy);
                Self::from_entropy(entropy)
            },
            Count::Words21 => {
                let mut entropy = [0u8; 28];
                rng.fill_bytes(&mut entropy);
                Self::from_entropy(entropy)
            },
            Count::Words24 => {
                let mut entropy = [0u8; 32];
                rng.fill_bytes(&mut entropy);
                Self::from_entropy(entropy)
            },
        }
        .expect("valid entropy length won't fail to generate the mnemonic")
    }

    /// Creates a new [`Mnemonic`] from the given entropy.
    ///
    /// # Example
    ///
    /// ```rust
    /// use bip0039::Mnemonic;
    ///
    /// let entropy = vec![0x1a, 0x48, 0x6a, 0x5f, 0xbe, 0x53, 0x63, 0x99, 0x84, 0xcb, 0x64, 0xb0, 0x70, 0x75, 0x5f, 0x7b];
    /// let mnemonic = <Mnemonic>::from_entropy(entropy).unwrap();
    /// assert_eq!(mnemonic.phrase(), "bottom drive obey lake curtain smoke basket hold race lonely fit walk");
    /// ```
    pub fn from_entropy<E: Into<Vec<u8>>>(entropy: E) -> Result<Self, Error> {
        let entropy = entropy.into();
        let phrase = encode_entropy::<L>(&entropy)?;

        Ok(Self {
            lang: PhantomData::<L>,
            phrase: Zeroizing::new(phrase),
            entropy: Zeroizing::new(entropy),
        })
    }

    /// Creates a [`Mnemonic`] from an existing mnemonic phrase.
    ///
    /// This method will normalize the input (UTF-8 NFKD) and normalize whitespace
    /// (single ASCII spaces).
    ///
    /// # Examples
    ///
    /// ```rust
    /// use bip0039::{Error, Mnemonic};
    ///
    /// let phrase = "bottom drive obey lake curtain smoke basket hold race lonely fit walk";
    /// let mnemonic = <Mnemonic>::from_phrase(phrase).unwrap();
    /// assert_eq!(mnemonic.phrase(), phrase);
    ///
    /// let phrase = "bottom drive obey lake curtain smoke basket hold race lonely fit shit";
    /// let mnemonic = <Mnemonic>::from_phrase(phrase);
    /// assert_eq!(mnemonic.unwrap_err(), Error::UnknownWord("shit".into()));
    /// ```
    pub fn from_phrase<'a, P: Into<Cow<'a, str>>>(phrase: P) -> Result<Self, Error> {
        let mut phrase = phrase.into();
        normalize_utf8(&mut phrase);

        Self::from_normalized_phrase(phrase)
    }

    /// Creates a [`Mnemonic`] from a phrase that is already normalized.
    ///
    /// Use this when you can guarantee the input is already normalized to UTF-8 NFKD.
    ///
    /// This avoids the UTF-8 NFKD normalization step during decoding, since the input is already
    /// normalized.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use bip0039::{Error, Mnemonic};
    ///
    /// let phrase = "bottom drive obey lake curtain smoke basket hold race lonely fit walk";
    /// let mnemonic = <Mnemonic>::from_normalized_phrase(phrase).unwrap();
    /// assert_eq!(mnemonic.phrase(), phrase);
    ///
    /// let phrase = "bottom drive obey lake curtain smoke basket hold race lonely fit shit";
    /// let mnemonic = <Mnemonic>::from_normalized_phrase(phrase);
    /// assert_eq!(mnemonic.unwrap_err(), Error::UnknownWord("shit".into()));
    /// ```
    pub fn from_normalized_phrase<'a, P: Into<Cow<'a, str>>>(phrase: P) -> Result<Self, Error> {
        let phrase = phrase.into();

        let decoded = decode_phrase::<L>(&phrase, DecodeMode::BuildNormalizedPhrase)?;
        let entropy = decoded.entropy;
        let normalized_phrase = decoded
            .normalized_phrase
            .expect("BuildNormalizedPhrase always constructs a normalized phrase");

        Ok(Mnemonic {
            lang: PhantomData::<L>,
            phrase: Zeroizing::new(normalized_phrase),
            entropy: Zeroizing::new(entropy),
        })
    }

    /// Validates the word count and checksum of a mnemonic phrase.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use bip0039::{Error, Mnemonic};
    /// use unicode_normalization::UnicodeNormalization;
    ///
    /// let phrase = "bottom drive obey lake curtain smoke basket hold race lonely fit walk";
    /// let result = <Mnemonic>::validate(phrase);
    /// assert!(result.is_ok());
    ///
    /// let phrase = "bottom drive obey lake curtain smoke basket hold race lonely fit shit";
    /// let result = <Mnemonic>::validate(phrase);
    /// assert_eq!(result.unwrap_err(), Error::UnknownWord("shit".into()));
    /// ```
    #[cfg_attr(
        feature = "japanese",
        doc = r##"
```rust
use bip0039::{Error, Japanese, Mnemonic};
use unicode_normalization::UnicodeNormalization;

let phrase = "そつう れきだい ほんやく わかす りくつ ばいか ろせん やちん そつう れきだい ほんやく わかめ";
let result = <Mnemonic<Japanese>>::validate(phrase);
assert!(result.is_ok());

let phrase = "そつう れきだい ほんやく わかす りくつ ばいか ろせん やちん そつう れきだい ほんやく ばか";
let result = <Mnemonic<Japanese>>::validate(phrase);
assert_eq!(result.unwrap_err(), Error::UnknownWord("ばか".nfkd().to_string()));
```
"##
    )]
    pub fn validate<'a, P: Into<Cow<'a, str>>>(phrase: P) -> Result<(), Error> {
        let mut phrase = phrase.into();
        normalize_utf8(&mut phrase);

        let _decoded = decode_phrase::<L>(&phrase, DecodeMode::ValidateOnly)?;
        Ok(())
    }

    /// Generates the seed from the [`Mnemonic`] and the passphrase.
    ///
    /// If a passphrase is not present, an empty string `""` is used instead.
    ///
    /// # Example
    ///
    /// ```rust
    /// use bip0039::Mnemonic;
    ///
    /// let phrase = "bottom drive obey lake curtain smoke basket hold race lonely fit walk";
    /// let mnemonic = <Mnemonic>::from_phrase(phrase).unwrap();
    /// assert_eq!(
    ///     mnemonic.to_seed("").to_vec(),
    ///     const_hex::decode("02d5cd1db85b4d1397d78978062a1160e76e94cc5aaad3089644846865bb18fc68ddf383059d3fe82902a203d60790a8c8ab488de5013d10a8a8bded8d9174b9").unwrap()
    /// );
    /// ```
    pub fn to_seed<P: AsRef<str>>(&self, passphrase: P) -> [u8; 64] {
        // use the PBKDF2 function with a mnemonic sentence (in UTF-8 NFKD) used as the password
        // and the string "mnemonic" + passphrase (again in UTF-8 NFKD) used as the salt.
        // The iteration count is set to 2048 and HMAC-SHA512 is used as the pseudo-random function.
        // The length of the derived key is 512 bits (= 64 bytes).
        const PBKDF2_ROUNDS: u32 = 2048;
        const PBKDF2_BYTES: usize = 64;

        // the phrase has been normalized
        let normalized_password = self.phrase();
        let normalized_salt = {
            let mut salt = Cow::Owned(format!("mnemonic{}", passphrase.as_ref()));
            normalize_utf8(&mut salt);
            salt
        };

        let mut seed = [0u8; PBKDF2_BYTES];
        pbkdf2::pbkdf2::<Hmac<Sha512>>(
            normalized_password.as_bytes(),
            normalized_salt.as_bytes(),
            PBKDF2_ROUNDS,
            &mut seed,
        )
        .expect("HMAC can be initialized with any key length");
        seed
    }

    /// Returns the mnemonic phrase as a string slice.
    pub fn phrase(&self) -> &str {
        &self.phrase
    }

    /// Consumes the `Mnemonic` and return the phrase as a `String`.
    pub fn into_phrase(mut self) -> String {
        // Create an empty string and swap values with the mnemonic's phrase.
        // This allows `Mnemonic` to implement `Drop`, while still returning the phrase.
        mem::take(&mut self.phrase)
    }

    /// Returns the original entropy of the mnemonic phrase.
    pub fn entropy(&self) -> &[u8] {
        &self.entropy
    }

    /// Consumes the `Mnemonic` and return the entropy as a `Vec<u8>`.
    pub fn into_entropy(mut self) -> Vec<u8> {
        // Create an empty bytes and swap values with the mnemonic's entropy.
        // This allows `Mnemonic` to implement `Drop`, while still returning the entropy.
        mem::take(&mut self.entropy)
    }
}

/// Ensure the content of the `s` is normalized UTF8.
/// Avoid allocation for normalization when there are no special UTF8 characters in the string.
#[inline]
fn normalize_utf8(s: &mut Cow<'_, str>) {
    if is_nfkd_quick(s.as_ref().chars()) != IsNormalized::Yes {
        *s = Cow::Owned(s.as_ref().nfkd().to_string())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_mnemonic_roundtrip() {
        let mnemonic = <Mnemonic>::generate(Count::Words12);
        let entropy = mnemonic.entropy();
        let phrase = mnemonic.phrase();

        {
            let m = <Mnemonic>::from_entropy(entropy).unwrap();
            assert_eq!(phrase, m.phrase());
        }

        {
            let m = <Mnemonic>::from_phrase(phrase).unwrap();
            assert_eq!(entropy, m.entropy());
        }
    }

    #[test]
    fn test_mnemonic_zeroize_when_drop() {
        let p: *const String;
        let e: *const Vec<u8>;

        // phrase = "absurd amount doctor acoustic avoid letter advice cage absurd amount doctor
        // adjust" entropy = [1u8; 16]
        {
            let m = <Mnemonic>::from_entropy([1u8; 16]).unwrap();
            p = &*m.phrase;
            e = &*m.entropy;
            unsafe {
                println!("*p: {}", *p);
                println!("*e: {:?}", *e);
            }
        }

        unsafe {
            assert_ne!(
                *p,
                "absurd amount doctor acoustic avoid letter advice cage absurd amount doctor adjust"
            );
            println!("*p: {}", *p);
            assert_ne!(*e, [1u8; 16]);
            println!("*e: {:?}", *e);
        }
    }

    #[test]
    fn test_mnemonic_consume() {
        let p: *const String;
        {
            let m = <Mnemonic>::from_entropy([1u8; 16]).unwrap();
            p = &*m.phrase;
            unsafe {
                println!("*p: {} ({:p})", *p, p);
            }
            let phrase = m.into_phrase();
            assert_ne!(p, &phrase);
            println!("phrase: {} ({:p})", phrase, &phrase);
        }
        // error
        // unsafe {
        //     println!("*p: {} ({:p})", (*p), p);
        // }

        let e: *const Vec<u8>;
        {
            let m = <Mnemonic>::from_entropy([1u8; 16]).unwrap();
            e = &*m.entropy;
            unsafe {
                println!("*e: {:?} ({:p})", *e, e);
            }
            let entropy = m.into_entropy();
            assert_ne!(e, &entropy);
            println!("entropy: {:?} ({:p})", entropy, &entropy);
        }
        // error
        // unsafe {
        //     println!("*e: {:?} ({:p})", (*e), e);
        // }
    }
}
