// Bitcoin Dev Kit
// Written in 2020 by Alekos Filini <alekos.filini@gmail.com>
//
// Copyright (c) 2020-2025 Bitcoin Dev Kit Developers
//
// This file is licensed under the Apache License, Version 2.0 <LICENSE-APACHE
// or http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your option.
// You may not use this file except in accordance with one or both of these
// licenses.

//! BIP-0039

// TODO: maybe write our own implementation of bip39? Seems stupid to have an extra dependency for
// something that should be fairly simple to re-implement.

use alloc::string::String;

use bitcoin::{bip32, Network};
use miniscript::ScriptContext;

use super::{
    any_network_kind, DerivableKey, DescriptorKey, ExtendedKey, GeneratableKey, GeneratedKey,
    KeyError,
};

pub use bip39::{Error, Language, Mnemonic};

type Seed = [u8; 64];

/// Type describing entropy length (aka word count) in the mnemonic
pub enum WordCount {
    /// 12 words mnemonic (128 bits entropy)
    Words12 = 128,
    /// 15 words mnemonic (160 bits entropy)
    Words15 = 160,
    /// 18 words mnemonic (192 bits entropy)
    Words18 = 192,
    /// 21 words mnemonic (224 bits entropy)
    Words21 = 224,
    /// 24 words mnemonic (256 bits entropy)
    Words24 = 256,
}

fn set_valid_on_any_network_kind<Ctx: ScriptContext>(
    descriptor_key: DescriptorKey<Ctx>,
) -> DescriptorKey<Ctx> {
    // We have to pick one network kind to build the xprv, but since the BIP39 standard doesn't
    // encode the network kind, the xprv we create is actually valid everywhere. So we override the
    // valid network kinds with `any_network_kind()`.
    descriptor_key.override_valid_network_kinds(any_network_kind())
}

/// Type for a BIP39 mnemonic with an optional passphrase.
pub type MnemonicWithPassphrase = (Mnemonic, Option<String>);

#[cfg_attr(docsrs, doc(cfg(feature = "keys-bip39")))]
impl<Ctx: ScriptContext> DerivableKey<Ctx> for Seed {
    fn into_extended_key(self) -> Result<ExtendedKey<Ctx>, KeyError> {
        Ok(bip32::Xpriv::new_master(Network::Bitcoin, &self[..])?.into())
    }

    fn into_descriptor_key(
        self,
        source: Option<bip32::KeySource>,
        derivation_path: bip32::DerivationPath,
    ) -> Result<DescriptorKey<Ctx>, KeyError> {
        let descriptor_key = self
            .into_extended_key()?
            .into_descriptor_key(source, derivation_path)?;

        Ok(set_valid_on_any_network_kind(descriptor_key))
    }
}

#[cfg_attr(docsrs, doc(cfg(feature = "keys-bip39")))]
impl<Ctx: ScriptContext> DerivableKey<Ctx> for MnemonicWithPassphrase {
    fn into_extended_key(self) -> Result<ExtendedKey<Ctx>, KeyError> {
        let (mnemonic, passphrase) = self;
        let seed: Seed = mnemonic.to_seed(passphrase.as_deref().unwrap_or(""));

        seed.into_extended_key()
    }

    fn into_descriptor_key(
        self,
        source: Option<bip32::KeySource>,
        derivation_path: bip32::DerivationPath,
    ) -> Result<DescriptorKey<Ctx>, KeyError> {
        let descriptor_key = self
            .into_extended_key()?
            .into_descriptor_key(source, derivation_path)?;

        Ok(set_valid_on_any_network_kind(descriptor_key))
    }
}

#[cfg_attr(docsrs, doc(cfg(feature = "keys-bip39")))]
impl<Ctx: ScriptContext> DerivableKey<Ctx> for (GeneratedKey<Mnemonic, Ctx>, Option<String>) {
    fn into_extended_key(self) -> Result<ExtendedKey<Ctx>, KeyError> {
        let (mnemonic, passphrase) = self;
        (mnemonic.into_key(), passphrase).into_extended_key()
    }

    fn into_descriptor_key(
        self,
        source: Option<bip32::KeySource>,
        derivation_path: bip32::DerivationPath,
    ) -> Result<DescriptorKey<Ctx>, KeyError> {
        let (mnemonic, passphrase) = self;
        (mnemonic.into_key(), passphrase).into_descriptor_key(source, derivation_path)
    }
}

#[cfg_attr(docsrs, doc(cfg(feature = "keys-bip39")))]
impl<Ctx: ScriptContext> DerivableKey<Ctx> for Mnemonic {
    fn into_extended_key(self) -> Result<ExtendedKey<Ctx>, KeyError> {
        (self, None).into_extended_key()
    }

    fn into_descriptor_key(
        self,
        source: Option<bip32::KeySource>,
        derivation_path: bip32::DerivationPath,
    ) -> Result<DescriptorKey<Ctx>, KeyError> {
        let descriptor_key = self
            .into_extended_key()?
            .into_descriptor_key(source, derivation_path)?;

        Ok(set_valid_on_any_network_kind(descriptor_key))
    }
}

#[cfg_attr(docsrs, doc(cfg(feature = "keys-bip39")))]
impl<Ctx: ScriptContext> GeneratableKey<Ctx> for Mnemonic {
    type Entropy = [u8; 32];

    type Options = (WordCount, Language);
    type Error = Option<bip39::Error>;

    fn generate_with_entropy(
        (word_count, language): Self::Options,
        entropy: Self::Entropy,
    ) -> Result<GeneratedKey<Self, Ctx>, Self::Error> {
        let entropy = &entropy[..(word_count as usize / 8)];
        let mnemonic = Mnemonic::from_entropy_in(language, entropy)?;

        Ok(GeneratedKey::new(mnemonic, any_network_kind()))
    }
}

#[cfg_attr(coverage_nightly, coverage(off))]
#[cfg(test)]
mod test {
    use super::WordCount;

    use alloc::string::ToString;
    use core::str::FromStr;

    use bip39::{Language, Mnemonic};
    use bitcoin::bip32;

    use crate::keys::{any_network_kind, GeneratableKey, GeneratedKey};

    #[test]
    fn test_keys_bip39_mnemonic() {
        let mnemonic =
            "aim bunker wash balance finish force paper analyst cabin spoon stable organ";
        let mnemonic = Mnemonic::parse_in(Language::English, mnemonic).unwrap();
        let path = bip32::DerivationPath::from_str("m/44'/0'/0'/0").unwrap();

        let key = (mnemonic, path);
        let (desc, keys, network_kinds) = crate::descriptor!(wpkh(key)).unwrap();
        assert_eq!(desc.to_string(), "wpkh([be83839f/44'/0'/0']xpub6DCQ1YcqvZtSwGWMrwHELPehjWV3f2MGZ69yBADTxFEUAoLwb5Mp5GniQK6tTp3AgbngVz9zEFbBJUPVnkG7LFYt8QMTfbrNqs6FNEwAPKA/0/*)#0r8v4nkv");
        assert_eq!(keys.len(), 1);
        assert_eq!(network_kinds, any_network_kind());
    }

    #[test]
    fn test_keys_bip39_mnemonic_passphrase() {
        let mnemonic =
            "aim bunker wash balance finish force paper analyst cabin spoon stable organ";
        let mnemonic = Mnemonic::parse_in(Language::English, mnemonic).unwrap();
        let path = bip32::DerivationPath::from_str("m/44'/0'/0'/0").unwrap();

        let key = ((mnemonic, Some("passphrase".into())), path);
        let (desc, keys, network_kinds) = crate::descriptor!(wpkh(key)).unwrap();
        assert_eq!(desc.to_string(), "wpkh([8f6cb80c/44'/0'/0']xpub6DWYS8bbihFevy29M4cbw4ZR3P5E12jB8R88gBDWCTCNpYiDHhYWNywrCF9VZQYagzPmsZpxXpytzSoxynyeFr4ZyzheVjnpLKuse4fiwZw/0/*)#h0j0tg5m");
        assert_eq!(keys.len(), 1);
        assert_eq!(network_kinds, any_network_kind());
    }

    #[test]
    fn test_keys_generate_bip39() {
        let generated_mnemonic: GeneratedKey<_, miniscript::Segwitv0> =
            Mnemonic::generate_with_entropy(
                (WordCount::Words12, Language::English),
                crate::keys::test::TEST_ENTROPY,
            )
            .unwrap();
        assert_eq!(generated_mnemonic.valid_network_kinds, any_network_kind());
        assert_eq!(
            generated_mnemonic.to_string(),
            "primary fetch primary fetch primary fetch primary fetch primary fetch primary fever"
        );

        let generated_mnemonic: GeneratedKey<_, miniscript::Segwitv0> =
            Mnemonic::generate_with_entropy(
                (WordCount::Words24, Language::English),
                crate::keys::test::TEST_ENTROPY,
            )
            .unwrap();
        assert_eq!(generated_mnemonic.valid_network_kinds, any_network_kind());
        assert_eq!(generated_mnemonic.to_string(), "primary fetch primary fetch primary fetch primary fetch primary fetch primary fetch primary fetch primary fetch primary fetch primary fetch primary fetch primary foster");
    }

    #[test]
    fn test_keys_generate_bip39_random() {
        let generated_mnemonic: GeneratedKey<_, miniscript::Segwitv0> =
            Mnemonic::generate((WordCount::Words12, Language::English)).unwrap();
        assert_eq!(generated_mnemonic.valid_network_kinds, any_network_kind());

        let generated_mnemonic: GeneratedKey<_, miniscript::Segwitv0> =
            Mnemonic::generate((WordCount::Words24, Language::English)).unwrap();
        assert_eq!(generated_mnemonic.valid_network_kinds, any_network_kind());
    }
}
