/*
 * 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/>.
 */
#include "HDMasterKey.h"
#include "HDMaster_p.h"

#include <crypto/hmac_sha512.h>
#include <crypto/pbkdf512.h>
#include "base58.h"
#include <cassert>

// static method
std::vector<uint32_t> HDMasterKey::deriveFromString(const std::string &path)
{
    if (path.empty() || path[0] != 'm')
        throw std::runtime_error("invalid path");
    /*
     * we expect a string like m/1/2'/3
     * where the 'm' indicates us as a starting point.
     * where each number is a derivation number.
     * where the single quote alters the number to be 'hardened', adding bit 32 to the number.
     */
    size_t slash = path.find('/');
    if (slash != 1) // the slash should follow the 'm'
        throw std::runtime_error("invalid path");
    std::vector<uint32_t> pathVector;
    // this reads numbers till we find an unknown char.
    while (path.size() > slash + 1 && path.at(slash) == '/') {
        size_t numChars = 0; // number of characters that are numbers.
        uint32_t num = std::stoi(path.substr(slash + 1), &numChars);
        if (numChars == 0)
            throw std::runtime_error("Derive:ERROR. Parsing error (expected number)");
        if (num >= Hardened)
            throw std::runtime_error("Derive:ERROR. Number out of range");
        if (path.size() > slash + 1 + numChars && path.at(slash + 1 + numChars) == '\'') {
            assert(num < Hardened);
            num += Hardened;
            ++numChars;
        }
        pathVector.push_back(num);
        slash += numChars + 1;
    }
    if (path.size() > slash // trailing chars
            && !(path.size() == slash + 1 && path.at(slash) == '/')) // oh, a single slash, thats Ok
        throw std::runtime_error("unrecognized trailing text");

    return pathVector;
}

uint8_t* HDMasters::prefix(Base58Type type)
{
    static uint8_t PubTestPrefix[4] = { 4, 0x35, 0x87, 0xCF };
    static uint8_t PubMainPrefix[4] = { 4, 0x88, 0xB2, 0x1E };
    static uint8_t PrivTestPrefix[4] = { 4, 0x35, 0x83, 0x94 };
    static uint8_t PrivMainPrefix[4] = { 4, 0x88, 0xAD, 0xE4 };

    switch (type) {
    case PublicKeyTestPrefix:
        return PubTestPrefix;
    case PublicKeyMainPrefix:
        return PubMainPrefix;
    case PrivateKeyTestPrefix:
        return PrivTestPrefix;
    case PrivateKeyMainPrefix:
        return PrivMainPrefix;
    }
    assert(false);
    return nullptr;
}


HDMasterKey::HDMasterKey()
{
    memset(m_fingerprint, 0, 4);
}

HDMasterKey HDMasterKey::fromMnemonic(const std::string &phrase, MnemonicType format, const std::string &password)
{
    const std::string salt = (format == ElectrumMnemonic ? "electrum" : "mnemonic") + password;
    std::vector<uint8_t> output(64);
    auto rc = pbkdf512(reinterpret_cast<const uint8_t*>(phrase.c_str()), phrase.size(),
                       reinterpret_cast<const uint8_t*>(salt.c_str()), salt.size(),
                       output.data(), output.size(), /* iterations = */ 2048);
    if (rc == 0)
        return fromSeed(output);
    return HDMasterKey();
}

HDMasterKey HDMasterKey::fromSeed(const std::vector<uint8_t> &seed)
{
    static const unsigned char hashkey[] = {'B','i','t','c','o','i','n',' ','s','e','e','d'};
    std::vector<uint8_t, secure_allocator<uint8_t>> out(64);
    CHMAC_SHA512(hashkey, sizeof(hashkey))
        .write(seed.data(), seed.size())
        .finalize(out.data());
    HDMasterKey answer;
    answer.m_key.set(out.data(), out.data() + 32, true);
    memcpy(answer.m_chaincode.begin(), out.data() + 32, 32);
    memset(answer.m_fingerprint, 0, sizeof(m_fingerprint)); // all zero's
    return answer;
}

HDMasterKey HDMasterKey::fromXPriv(const std::string &xpriv, Chain *chain)
{
    std::vector<uint8_t> data;
    const bool ok = Base58::decodeCheck(xpriv.c_str(), data);
    if (!ok || data.size() != 78)
        throw std::runtime_error("invalid xpriv");

    auto prefix = HDMasters::prefix(HDMasters::PrivateKeyMainPrefix);
    if (memcmp(prefix, data.data(), 4) == 0) {
        if (chain)
            *chain = MainChain;
    } else {
        prefix = HDMasters::prefix(HDMasters::PrivateKeyTestPrefix);
        if (memcmp(prefix, data.data(), 4) != 0)
            throw std::runtime_error("Not a recognized xpriv chain");
        if (chain)
            *chain = Testnet;
    }

    HDMasterKey answer;
    uint8_t *code = data.data() + 4;
    answer.m_level = code[0];
    static_assert(sizeof(HDMasterKey::m_fingerprint) == 4);
    memcpy(answer.m_fingerprint, code + 1, 4);
    answer.m_childId = code[5] << 24;
    answer.m_childId += code[6] << 16;
    answer.m_childId += code[7] << 8;
    answer.m_childId += code[8];
    memcpy(&answer.m_chaincode, code+9, 32);
    answer.m_key.set(code + 42, code + 42 + 32);
    return answer;
}

bool HDMasterKey::isValid() const
{
    return m_key.isValid();
}

std::string HDMasterKey::serialize(const std::vector<uint32_t> &path, HDMasterKey::Chain chain, HDMasterKey::SerializeType type) const
{
    assert(path.size() < 256);
    std::vector<uint8_t> data(78);
    if (type == _XPriv)
        memcpy(data.data(), HDMasters::prefix((chain == MainChain) ? HDMasters::PrivateKeyMainPrefix : HDMasters::PrivateKeyTestPrefix), 4);
    else
        memcpy(data.data(), HDMasters::prefix((chain == MainChain) ? HDMasters::PublicKeyMainPrefix : HDMasters::PublicKeyTestPrefix), 4);
    // derive it
    assert(m_key.isValid());
    assert(m_key.isCompressed());
    ChainCode currentChaincode = m_chaincode;
    PrivateKey currentKey = m_key;

    KeyId parentFingerPrint;
    for (const uint32_t child : path) {
        PrivateKey newKey;
        ChainCode newChaincode;
        bool ok = currentKey.derive(newKey, newChaincode, child, currentChaincode);
        if (!ok)
            return std::string();
        parentFingerPrint = currentKey.getPubKey().getKeyId();
        currentKey = newKey;
        currentChaincode = newChaincode;
    }

    uint8_t *code = data.data() + 4;
    code[0] = path.size(); // depth
    memcpy(code+1, &parentFingerPrint, 4);
    uint32_t nChild = path.empty() ? 0 : path.back();
    code[5] = (nChild >> 24) & 0xFF;
    code[6] = (nChild >> 16) & 0xFF;
    code[7] = (nChild >>  8) & 0xFF;
    code[8] = (nChild >>  0) & 0xFF;
    memcpy(code+9, currentChaincode.begin(), 32);

    if (type == _XPriv) {
        code[41] = 0;
        assert(m_key.size() == 32);
        memcpy(code+42, currentKey.begin(), 32);
    }
    else {
        auto pubkey = currentKey.getPubKey();
        assert(pubkey.size() == 33);
        memcpy(code+41, pubkey.begin(), 33);
    }
    return Base58::encodeCheck(data);
}

int HDMasterKey::level() const
{
    return m_level;
}

std::string HDMasterKey::toXPubString(const std::vector<uint32_t> &path, Chain chain) const
{
    return serialize(path, chain, _XPub);
}

std::string HDMasterKey::toXPrivString(const std::vector<uint32_t> &path, Chain chain) const
{
    return serialize(path, chain, _XPriv);
}

PrivateKey HDMasterKey::derive(const std::vector<uint32_t> &path) const
{
    assert(m_key.isValid());
    assert(m_key.isCompressed());
    ChainCode currentChaincode = m_chaincode;
    PrivateKey currentKey = m_key;

    int level = m_level;
    for (const uint32_t child : path) {
        if (level-- > 0) {
            if (level == 0 && child != m_childId)
                throw std::runtime_error("Deriving on different path than this sub-key is at");
            continue;
        }
        PrivateKey newKey;
        ChainCode newChaincode;
        bool ok = currentKey.derive(newKey, newChaincode, child, currentChaincode);
        if (!ok)
            return PrivateKey(); // return invalid key
        currentKey = newKey;
        currentChaincode = newChaincode;
    }
    return currentKey;
}

PrivateKey HDMasterKey::derive(const std::string &path) const
{
    return derive(HDMasterKey::deriveFromString(path));
}