/*
 * This file is part of the Flowee project
 * Copyright (C) 2021 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 <crypto/hmac_sha512.h>
#include <crypto/pbkdf512.h>
#include "base58.h"

// 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;
}


HDMasterKey::HDMasterKey()
{
}

HDMasterKey HDMasterKey::fromMnemonic(const std::string &phrase, const std::string &password)
{
    std::string pwd = "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*>(pwd.c_str()), pwd.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));
    return answer;
}

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

std::string HDMasterKey::toString(Chain chain) const
{
    static uint8_t TestPrefix[4] = { 4, 0x35, 0x87, 0xCF };
    static uint8_t MainPrefix[4] = { 4, 0x88, 0x82, 0x1E };
    std::vector<uint8_t> data(78);
    memcpy(data.data(),(chain == MainChain) ? MainPrefix : TestPrefix, 4);
    uint8_t *code = data.data() + 4;
    auto pubkey = m_key.GetPubKey();
    code[0] = 0; // depth
    memcpy(code+1, m_fingerprint, 4);
    uint32_t nChild = 0; // master, this is always zero
    code[5] = (nChild >> 24) & 0xFF; code[6] = (nChild >> 16) & 0xFF;
    code[7] = (nChild >>  8) & 0xFF; code[8] = (nChild >>  0) & 0xFF;
    memcpy(code+9, m_chaincode.begin(), 32);
    assert(pubkey.size() == 33);
    memcpy(code+41, pubkey.begin(), 33);
    return EncodeBase58Check(data);
}

std::string HDMasterKey::privToString(Chain chain) const
{
    static uint8_t TestPrefix[4] = { 4, 0x35, 0x83, 0x94 };
    static uint8_t MainPrefix[4] = { 4, 0x88, 0xAD, 0xE4 };
    std::vector<uint8_t> data(78);
    memcpy(data.data(),(chain == MainChain) ? MainPrefix : TestPrefix, 4);

    uint8_t *code = data.data() + 4;
    code[0] = 0; // depth
    memcpy(code+1, m_fingerprint, 4);
    uint32_t nChild = 0; // master, this is always zero
    code[5] = (nChild >> 24) & 0xFF; code[6] = (nChild >> 16) & 0xFF;
    code[7] = (nChild >>  8) & 0xFF; code[8] = (nChild >>  0) & 0xFF;
    memcpy(code+9, m_chaincode.begin(), 32);
    code[41] = 0;
    assert(m_key.size() == 32);
    memcpy(code+42, m_key.begin(), 32);
    return EncodeBase58Check(data);
}

CKey HDMasterKey::derive(const std::vector<uint32_t> &path) const
{
    assert(m_key.IsValid());
    assert(m_key.IsCompressed());
    ChainCode currentChaincode = m_chaincode;
    CKey currentKey = m_key;

    for (const uint32_t child : path) {
        CKey newKey;
        ChainCode newChaincode;
        bool ok = currentKey.Derive(newKey, newChaincode, child, currentChaincode);
        if (!ok)
            return CKey(); // return invalid key
        currentKey = newKey;
        currentChaincode = newChaincode;
    }
    return currentKey;
}

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