/*
 * This file is part of the Flowee project
 * Copyright (C) 2021-2022 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 "HDMasterPubkey.h"
#include "HDMaster_p.h"
#include "HDMasterKey.h"

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

// static helper method
std::vector<uint32_t> HDMasterPubkey::deriveFromString(const std::string &path)
{
    return HDMasterKey::deriveFromString(path);
}

HDMasterPubkey HDMasterPubkey::fromXPub(const std::string &xpub)
{
    std::vector<uint8_t> data;
    const bool ok = DecodeBase58Check(xpub.c_str(), data);
    if (!ok || data.size() != 78)
        throw std::runtime_error("invalid xpub");
    HDMasterPubkey answer;
    if (memcmp(HDMasters::prefix(HDMasters::PublicKeyTestPrefix), data.data(), 4) == 0)
        answer.m_chain = Testnet;
    else if (memcmp(HDMasters::prefix(HDMasters::PublicKeyMainPrefix), data.data(), 4) == 0)
        answer.m_chain = MainChain;
    else
        throw std::runtime_error("Invalid prefix");

    const uint8_t *code = data.data() + 4;
    answer.m_depth = code[0];
    answer.m_lastChild = code[8] + (code[7] << 8) + (code[6] << 16) + (code[5] << 24);
    memcpy(answer.m_fingerprint, code+1, 4);
    memcpy(answer.m_chaincode.begin(), code+9, 32);
    answer.m_pubkey.setData(code+41, code+41 + 33);
    return answer;
}

HDMasterPubkey HDMasterPubkey::fromHDMaster(const HDMasterKey &key, const std::vector<uint32_t> &derivationPath_)
{
    // create a derivation path that limits itself to the hardened section.
    size_t relevantNum = derivationPath_.size();
    while (relevantNum > 0) {
        if (derivationPath_[relevantNum - 1]  & HDMasterKey::Hardened)
            break;
        --relevantNum;
    }
    std::vector<uint32_t> derivationPath;
    for (size_t i = 0; i < relevantNum; ++i) {
        derivationPath.push_back(derivationPath_[i]);
    }

    return HDMasterPubkey::fromXPub(key.toXPubString(derivationPath, HDMasterKey::MainChain));
}

HDMasterPubkey HDMasterPubkey::fromHDMaster(const HDMasterKey &key, const std::string &derivationPath)
{
    return fromHDMaster(key, HDMasterKey::deriveFromString(derivationPath));
}

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

HDMasterPubkey::Chain HDMasterPubkey::chain() const
{
    return m_chain;
}

bool HDMasterPubkey::isValid() const
{
    return m_pubkey.isValid();
}

std::string HDMasterPubkey::toString() const
{
    std::vector<uint8_t> data(78);
    memcpy(data.data(), HDMasters::prefix((m_chain == MainChain) ? HDMasters::PublicKeyMainPrefix : HDMasters::PublicKeyTestPrefix), 4);
    uint8_t *code = data.data() + 4;
    code[0] = m_depth;
    memcpy(code+1, m_fingerprint, 4);
    code[5] = (m_lastChild >> 24) & 0xFF; code[6] = (m_lastChild >> 16) & 0xFF;
    code[7] = (m_lastChild >>  8) & 0xFF; code[8] = (m_lastChild >>  0) & 0xFF;
    memcpy(code+9, m_chaincode.begin(), 32);
    assert(m_pubkey.size() == 33);
    memcpy(code+41, m_pubkey.begin(), 33);
    return EncodeBase58Check(data);
}

PublicKey HDMasterPubkey::derive(const std::vector<uint32_t> &path) const
{
    assert(m_pubkey.isValid());
    assert(m_pubkey.isCompressed());
    ChainCode currentChaincode = m_chaincode;
    PublicKey currentKey = m_pubkey;

    int levelsToSkip = m_depth;
    for (const uint32_t child : path) {
        if (--levelsToSkip >= 0)
            continue;
        PublicKey newKey;
        ChainCode newChaincode;
        bool ok = currentKey.derive(newKey, newChaincode, child, currentChaincode);
        if (!ok)
            return PublicKey(); // return invalid key
        currentKey = newKey;
        currentChaincode = newChaincode;
    }
    return currentKey;
}

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