use std::{error::Error, fmt::Display, str::FromStr, sync::Arc};

use hmac::Hmac;
use pbkdf2::pbkdf2;
use sha2::{Digest, Sha256, Sha512};

/// The error type representing a failure to create a [`Mnemonic`]. These errors only occur during
/// [`Mnemonic`] creation.
#[derive(Debug, Clone)]
pub enum MnemonicGenerationError {
    /// The amount of bits passed to a mnemonic must be divisible by 32.
    InvalidByteCount(usize),

    /// The length of a mnemonic in bits must be within the BIP-0039 range, and supported by the
    /// library. Currently, only 128, 192 (for testing purposes), and 256 are supported.
    InvalidByteLength(usize),

    /// Invalid length resulting from PBKDF2.
    InvalidPbkdf2Length,
}

impl Display for MnemonicGenerationError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            MnemonicGenerationError::InvalidByteCount(count) => {
                write!(f, "Invalid byte count: {count}, must be divisible by 8")
            }
            MnemonicGenerationError::InvalidByteLength(count) => {
                write!(f, "Invalid byte length: {count}, must be 128 or 256")
            }
            MnemonicGenerationError::InvalidPbkdf2Length => {
                f.write_str("Invalid length from PBKDF2")
            }
        }
    }
}

impl Error for MnemonicGenerationError {}

/// A BIP-0039 compatible list of words.
#[derive(Debug, Clone)]
pub struct Wordlist(Vec<String>);

impl Default for Wordlist {
    /// Returns the English wordlist in the Bitcoin BIP-0039 specification.
    fn default() -> Self {
        // TODO: English is the only supported language.
        let wordlist_file = include_str!("data/wordlist.txt");
        Wordlist(
            wordlist_file
                .lines()
                // skip 1: comment at top of file to point to BIP-0039 source.
                .skip(1)
                .map(|x| x.trim().to_string())
                .collect(),
        )
    }
}

impl Wordlist {
    /// Return an Arced version of the Wordlist
    #[allow(clippy::must_use_candidate)]
    pub fn arc(self) -> Arc<Self> {
        Arc::new(self)
    }

    /// Determine whether the Wordlist contains a given word.
    pub fn contains(&self, word: &str) -> bool {
        self.0.iter().any(|w| w.as_str() == word)
    }

    /// Given an index, get a word from the wordlist.
    pub fn get_word(&self, word: usize) -> Option<&String> {
        self.0.get(word)
    }

    /*
    fn inner(&self) -> &Vec<String> {
        &self.0
    }
    */

    #[cfg(test)]
    fn into_inner(self) -> Vec<String> {
        self.0
    }
}

/// A BIP-0039 mnemonic with reference to a [`Wordlist`].
#[derive(Debug, Clone)]
pub struct Mnemonic {
    entropy: Vec<u8>,
    // words: Vec<usize>,
    wordlist: Arc<Wordlist>,
}

impl Display for Mnemonic {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let bit_count = self.entropy.len() * 8;
        let mut bits = vec![false; bit_count + bit_count / 32];

        for byte_index in 0..bit_count / 8 {
            for bit_index in 0..8 {
                bits[byte_index * 8 + bit_index] =
                    (self.entropy[byte_index] & (1 << (7 - bit_index))) > 0;
            }
        }

        let mut hasher = Sha256::new();
        hasher.update(&self.entropy);
        let hash = hasher.finalize().to_vec();
        for check_bit in 0..bit_count / 32 {
            bits[bit_count + check_bit] = (hash[check_bit / 8] & (1 << (7 - (check_bit % 8)))) > 0;
        }

        let mut iter = bits
            .chunks_exact(11)
            .peekable()
            .map(|chunk| {
                let mut num = 0usize;
                for i in 0..11 {
                    num += usize::from(chunk[10 - i]) << i;
                }
                num
            })
            .filter_map(|word| self.wordlist.get_word(word))
            .peekable();
        while let Some(word) = iter.next() {
            f.write_str(&word)?;
            if iter.peek().is_some() {
                f.write_str(" ")?;
            }
        }
        Ok(())
    }
}

/// The error type representing a failure to parse a [`Mnemonic`]. These errors only occur during
/// [`Mnemonic`] creation.
#[derive(Debug, Clone)]
pub enum MnemonicFromStrError {
    /// The amount of words used to parse a mnemonic was not correct.
    InvalidWordCount(usize),

    /// One of the words used to generate the mnemonic was not found in the default wordlist.
    InvalidWord(usize),

    /// The checksum for the mnemonic did not match the given words.
    InvalidChecksum,
}

impl Display for MnemonicFromStrError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str("Mnemonic error: ")?;
        match self {
            MnemonicFromStrError::InvalidWordCount(count) => {
                write!(f, "Incorrect word count: {count}")
            }
            MnemonicFromStrError::InvalidWord(index) => {
                write!(f, "Unknown word at index: {index}")
            }
            MnemonicFromStrError::InvalidChecksum => {
                f.write_str("Checksum of data did not match expected value")
            }
        }
    }
}

impl Error for MnemonicFromStrError {}

impl FromStr for Mnemonic {
    type Err = MnemonicFromStrError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let words: Vec<_> = s.split_whitespace().collect();
        let mut usize_words = vec![];
        let wordlist = Wordlist::default().arc();
        let mut bits = vec![false; words.len() * 11];
        for (index, word) in words.iter().enumerate() {
            let word = wordlist
                .0
                .iter()
                .position(|w| &w == word)
                .ok_or(MnemonicFromStrError::InvalidWord(index))?;
            usize_words.push(word);
            for bit in 0..11 {
                bits[index * 11 + bit] = (word & (1 << (10 - bit))) > 0;
            }
        }

        let mut checksum_bits = vec![false; bits.len() - (bits.len() * 32 / 33)];
        checksum_bits.copy_from_slice(&bits[bits.len() * 32 / 33..]);

        // remove checksum bits
        bits.truncate(bits.len() * 32 / 33);
        // bits.truncate(bits.len() - bits.len() % 32);

        let entropy: Vec<u8> = bits
            .chunks_exact(8)
            .map(|chunk| {
                let mut num = 0u8;
                for i in 0..8 {
                    num += u8::from(chunk[7 - i]) << i;
                }
                num
            })
            .collect();

        let mut hasher = Sha256::new();
        hasher.update(&entropy);
        let hash = hasher.finalize().to_vec();

        for (i, bit) in checksum_bits.iter().enumerate() {
            if !hash[i / 8] & (1 << (7 - (i % 8))) == *bit as u8 {
                return Err(MnemonicFromStrError::InvalidChecksum);
            }
        }

        Ok(Mnemonic {
            entropy,
            // words: usize_words,
            wordlist,
        })
    }
}

impl Mnemonic {
    /// Generate a [`Mnemonic`] from the provided entropy and [`Wordlist`].
    ///
    /// # Errors
    /// An error may be returned if the entropy is not within the acceptable lengths.
    pub fn from_entropy(
        bytes: &[u8],
        wordlist: Arc<Wordlist>,
    ) -> Result<Mnemonic, MnemonicGenerationError> {
        let bit_count = bytes.len() * 8;

        if bit_count % 32 != 0 {
            return Err(MnemonicGenerationError::InvalidByteCount(bit_count));
        }
        // 192 supported for test suite
        if ![128, 192, 256].contains(&bit_count) {
            return Err(MnemonicGenerationError::InvalidByteLength(bit_count));
        }

        Ok(unsafe { Self::from_raw_entropy(bytes, wordlist) })
    }

    pub unsafe fn from_raw_entropy(bytes: &[u8], wordlist: Arc<Wordlist>) -> Mnemonic {
        Mnemonic {
            entropy: bytes.to_vec(),
            wordlist,
        }
    }

    pub fn as_bytes(&self) -> &[u8] {
        &self.entropy
    }

    pub fn to_bytes(self) -> Vec<u8> {
        self.entropy
    }

    pub fn entropy(&self) -> Vec<u8> {
        self.entropy.clone()
    }

    pub fn seed<'a>(
        &self,
        passphrase: impl Into<Option<&'a str>>,
    ) -> Result<Vec<u8>, MnemonicGenerationError> {
        let passphrase = passphrase.into();

        let mut seed = [0u8; 64];
        let mnemonic = self.to_string();
        let salt = ["mnemonic", passphrase.unwrap_or("")].join("");
        pbkdf2::<Hmac<Sha512>>(mnemonic.as_bytes(), salt.as_bytes(), 2048, &mut seed)
            .map_err(|_| MnemonicGenerationError::InvalidPbkdf2Length)?;
        Ok(seed.to_vec())
    }

    pub fn words(self) -> (Vec<usize>, Arc<Wordlist>) {
        let bit_count = self.entropy.len() * 8;
        let mut bits = vec![false; bit_count + bit_count / 32];

        for byte_index in 0..bit_count / 8 {
            for bit_index in 0..8 {
                bits[byte_index * 8 + bit_index] =
                    (self.entropy[byte_index] & (1 << (7 - bit_index))) > 0;
            }
        }

        let mut hasher = Sha256::new();
        hasher.update(&self.entropy);
        let hash = hasher.finalize().to_vec();
        for check_bit in 0..bit_count / 32 {
            bits[bit_count + check_bit] = (hash[check_bit / 8] & (1 << (7 - (check_bit % 8)))) > 0;
        }

        let words = bits.chunks_exact(11).peekable().map(|chunk| {
            let mut num = 0usize;
            for i in 0..11 {
                num += usize::from(chunk[10 - i]) << i;
            }
            num
        });
        (words.collect(), self.wordlist.clone())
    }
}

#[cfg(test)]
mod tests {
    use std::{collections::HashSet, fs::File, io::Read};

    use super::*;

    #[test]
    fn wordlist_word_count_correct() {
        let wordlist = Wordlist::default().into_inner();
        assert_eq!(
            wordlist.len(),
            2usize.pow(11),
            "Wordlist did not include correct word count"
        );
    }

    #[test]
    fn reproduces_its_own_seed() {
        let mut random_handle = File::open("/dev/random").unwrap();
        let entropy = &mut [0u8; 256 / 8];
        random_handle.read_exact(&mut entropy[..]).unwrap();
        let wordlist = Wordlist::default().arc();
        let mnemonic = super::Mnemonic::from_entropy(&entropy[..256 / 8], wordlist).unwrap();
        let new_entropy = mnemonic.entropy();
        assert_eq!(&new_entropy, entropy);
    }

    #[test]
    fn conforms_to_trezor_tests() {
        let content = include_str!("data/vectors.json");
        let jsonobj: serde_json::Value = serde_json::from_str(content).unwrap();
        let wordlist = Wordlist::default().arc();

        for test in jsonobj["english"].as_array().unwrap() {
            let [ref hex_, ref seed, ..] = test.as_array().unwrap()[..] else {
                panic!("bad test: {test}");
            };
            let hex = hex::decode(hex_.as_str().unwrap()).unwrap();

            let mnemonic = Mnemonic::from_entropy(&hex, wordlist.clone()).unwrap();

            assert_eq!(mnemonic.to_string(), seed.as_str().unwrap());
        }
    }

    #[test]
    fn matches_bip39_crate() {
        let mut random_handle = File::open("/dev/random").unwrap();
        let entropy = &mut [0u8; 256 / 8];
        random_handle.read_exact(&mut entropy[..]).unwrap();
        let wordlist = Wordlist::default().arc();
        let my_mnemonic = super::Mnemonic::from_entropy(&entropy[..256 / 8], wordlist).unwrap();
        let their_mnemonic = bip39::Mnemonic::from_entropy(&entropy[..256 / 8]).unwrap();
        assert_eq!(my_mnemonic.to_string(), their_mnemonic.to_string());
        assert_eq!(my_mnemonic.seed(None).unwrap(), their_mnemonic.to_seed(""));
        assert_eq!(
            my_mnemonic.seed("testing").unwrap(),
            their_mnemonic.to_seed("testing")
        );
        assert_ne!(
            my_mnemonic.seed("test1").unwrap(),
            their_mnemonic.to_seed("test2")
        );
    }

    #[test]
    fn count_rate_of_duplicate_words() {
        let tests = 100_000;
        let mut count = 0.;
        let entropy = &mut [0u8; 256 / 8];
        let wordlist = Wordlist::default().arc();
        let mut random = std::fs::File::open("/dev/urandom").unwrap();
        let mut hs = HashSet::<usize>::with_capacity(24);

        for _ in 0..tests {
            random.read_exact(&mut entropy[..]).unwrap();
            let mnemonic = Mnemonic::from_entropy(&entropy[..256 / 8], wordlist.clone()).unwrap();
            let (words, _) = mnemonic.words();
            hs.clear();
            hs.extend(words);
            if hs.len() != 24 {
                count += 1.;
            }
        }

        // NOTE: Birthday problem math is: 0.126532
        // Set values to (about) 1 below, 1 above
        // Source: https://en.wikipedia.org/wiki/Birthday_problem
        let min = 11.5;
        let max = 13.5;
        assert!(
            count > f64::from(tests) * min / 100.,
            "{count} probability should be more than {min}%: {}",
            count / f64::from(tests)
        );
        assert!(
            count < f64::from(tests) * max / 100.,
            "{count} probability should be more than {max}%: {}",
            count / f64::from(tests)
        );
    }

    #[test]
    fn can_do_up_to_1024_bits() {
        let entropy = &mut [0u8; 128];
        let wordlist = Wordlist::default().arc();
        let mut random = std::fs::File::open("/dev/urandom").unwrap();
        random.read_exact(&mut entropy[..]).unwrap();
        let mnemonic = unsafe { Mnemonic::from_raw_entropy(&entropy[..], wordlist.clone()) };
        let (words, _) = mnemonic.words();
        assert!(words.len() == 96);
    }
}