/*
 * This file is part of the Flowee project
 * Copyright (C) 2021-2024 Tom Zander <tom@flowee.org>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
#ifndef FLOWEE_HDMASTERKEY_H
#define FLOWEE_HDMASTERKEY_H

#include "primitives/PrivateKey.h"
#include <string>

/**
 * Hierarchically Deterministic wallet Master key.
 * This class stores a HD wallet masterkey and allow you to derive any child key
 * from it.
 * This implements BIP32 and combined with the fromMnemonic() method this supports BIP39 based wallets,
 * and can also read Electrum format seed phrases as well.
 */
class HDMasterKey
{
public:
    /// Adding this value to a path-element makes it 'hardened'
    static const uint32_t Hardened = 0x80000000;

    /// creates an invalid masterkey.
    HDMasterKey();
    HDMasterKey(const HDMasterKey &other) = default;

    /**
     * Convert a string based derivation path into a vector of fields.
     * This throws an exception when the string is not a fully valid derivation path.
     *
     * The minimum correct string is: "m/"
     */
    static std::vector<uint32_t> deriveFromString(const std::string &path);


    // Note these exact values get serialized to the wallet file in Flowee Pay so don't modify their values
    enum MnemonicType {
        BIP39Mnemonic = 0,      ///< The preferred mnemonic type (industry standard)
        ElectrumMnemonic = 1    ///< Support for legacy Electron Cash and Electrum wallets
    };

    /**
     * Used in the string output methods.
     */
    enum Chain {
        Testnet,
        MainChain
    };

    /**
     * Create a masterkey from a known-good mnemonic.
     * @see Mnemonic::validateMnemonic() to make sure the string is good.
     * @param phrase a UTF8-encoded multi-word mnemonic phrase
     * @param format specifies how to interpret this mnemonic phrase; note that BIP39Menomic is the industry standard
     *        and ElectrumMnemonic is a legacy format that this library can read (but cannot generate)
     */
    static HDMasterKey fromMnemonic(const std::string &phrase, MnemonicType format, const std::string &password = {});
    static HDMasterKey fromSeed(const std::vector<uint8_t> &seed);
    /**
     * create a masterkey instance from a xpriv string.
     * By passing in a chain argument, this will be used to set the chain this xpriv was encoded for.
     */
    static HDMasterKey fromXPriv(const std::string &xpriv, Chain *chain = nullptr);

    /// returns true if the key is valid.
    bool isValid() const;
    /// output xprv private key as text
    std::string toXPrivString(const std::vector<uint32_t> &path, Chain chain = MainChain) const;
    std::string toXPrivString(const std::string &path, Chain chain = MainChain) const {
        return toXPrivString(deriveFromString(path), chain);
    }
    /**
     * Create an xpub which can be used to create pub keys.
     *
     * To understand this method, one must first understand the concept of 'Hardened', as
     * used by this class.
     * A derivation can be something like "m/0'/2/10", which, when passed to derive()
     * generates a single private key. The single quote behind the zero indicates that
     * that level is Hardened. Which has the effect that only private keys can derive
     * that level.
     * This is a way to protect the master private key in case of partial leakage.
     *
     * The effect is that if we want to have a xpub that is actually able to compute
     * child public keys from a derivation path, we have to apply hardened derivations
     * before making the xpub.
     *
     * In the case of "m/0'/2/10" we can return an xpub for "m/0'/" and then using that xpub
     * it is posible to create any public key under that master pubkey.
     *
     * @see HDMasterPubkey
     *
     * @path path to derive first.
     * @chain which chain this xpub is for.
     */
    std::string toXPubString(const std::vector<uint32_t> &path, Chain chain = MainChain) const;
    std::string toXPubString(const std::string &path, Chain chain = MainChain) const {
        return toXPubString(deriveFromString(path), chain);
    }

    // notice that this may throw if deriving an already derived masterkey with incompatible path
    PrivateKey derive(const std::vector<uint32_t> &path) const;
    // notice that this may throw if deriving an already derived masterkey with incompatible path
    PrivateKey derive(const std::string &path) const;

    HDMasterKey &operator=(const HDMasterKey &other) = default;
    friend bool operator==(const HDMasterKey &a, const HDMasterKey &b) {
        return memcmp(&a.m_fingerprint[0], &b.m_fingerprint[0], 4) == 0 && a.m_chaincode == b.m_chaincode && a.m_key == b.m_key;
    }

    int level() const;

private:
    enum SerializeType { _XPriv, _XPub };
    std::string serialize(const std::vector<uint32_t> &path, Chain chain, SerializeType type) const;

    uint8_t m_fingerprint[4]; // the parent's key fingerprint
    ChainCode m_chaincode;
    PrivateKey m_key;
    int m_level = 0;
    /*
     * A masterkey can be pre-derived. For instance to m/10/11
     * in this example m_level is 2 and m_child is 11.
     */
    uint32_t m_childId = 0;
};

#endif