/*
 * This file is part of the Flowee project
 * Copyright (C) 2025 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 "WalletData.h"
#include "BackupSyncModuleInfo.h"

#include <AccountConfig.h>
#include <FloweePay.h>

#include <streaming/BufferPools.h>
#include <streaming/MessageBuilder.h>
#include <sha256.h>

#include <QFile>
#include <QTimer>
#include <QJsonDocument>
#include <QJsonObject>

#include <openssl/evp.h>
#include <openssl/kdf.h>
#include <openssl/err.h>
#include <openssl/params.h>
#include <openssl/core_names.h>

namespace BCHMDF {
enum Tags {
    BlockHeight = 1,    // int: height
    TransactionIdShort, // 64 bit hash
    TransactionComment, // utf-8 string
    LastBlockSeen       // int: height.
};
}

namespace {

std::vector<unsigned char> derivePrivKey(const std::vector<unsigned char> &input)
{
    assert (input.size() == 32);
    std::vector<uint8_t> answer;

    std::unique_ptr<EVP_KDF, decltype(&::EVP_KDF_free)> kdf(EVP_KDF_fetch(nullptr, "HKDF", nullptr), ::EVP_KDF_free);
    if (!kdf.get())
        return answer;
    std::unique_ptr<EVP_KDF_CTX, decltype(&::EVP_KDF_CTX_free)> kctx(EVP_KDF_CTX_new(kdf.get()), ::EVP_KDF_CTX_free);
    if (!kctx.get())
        return answer;

    OSSL_PARAM params[5];
    params[0] = OSSL_PARAM_construct_utf8_string(OSSL_KDF_PARAM_DIGEST, const_cast<char*>("SHA256"), 0);
    params[1] = OSSL_PARAM_construct_octet_string(OSSL_KDF_PARAM_KEY, const_cast<unsigned char*>(input.data()), input.size());
    std::vector<unsigned char> salt = { 'f', 'l', 'o', 'w', 'e', 'e', 'p', 'a', 'y' };
    uint8_t *info = nullptr;
    params[2] = OSSL_PARAM_construct_octet_string(OSSL_KDF_PARAM_SALT, const_cast<unsigned char*>(salt.data()), salt.size());
    params[3] = OSSL_PARAM_construct_octet_string(OSSL_KDF_PARAM_INFO, info, 0);
    params[4] = OSSL_PARAM_construct_end();

    if (EVP_KDF_CTX_set_params(kctx.get(), params) <= 0)
        return answer;
    answer.resize(32);
    if (EVP_KDF_derive(kctx.get(), answer.data(), 32, nullptr) <= 0) // Derive the key
        answer.resize(0); // failed. Return empty.
    return answer;
}

// Generate 12-byte nonce from given timestamp
std::vector<uint8_t> nonceFromTime(uint32_t timestamp)
{
    std::vector<uint8_t> nonce(12, 0);
    constexpr const char *padding = "BCH-MDF"; // Bitcoin Cash MetaData Format
    std::memcpy(nonce.data(), padding, 7);
    const uint8_t *timestamp_p = reinterpret_cast<const uint8_t*>(&timestamp);
    // Notice we leave 1 zero byte to allow the timestamp to reach 5 bytes.
    for (int i = 0; i < 4; ++i) // copy at end of our bytearray
        nonce[7 + i] = timestamp_p[i];
    return nonce;
}

Streaming::ConstBuffer encrypt(std::vector<uint8_t> &key, uint32_t timestamp, const Streaming::ConstBuffer &clearText, const std::shared_ptr<Streaming::BufferPool> &pool)
{
    assert(key.size() == 32);

    std::unique_ptr<EVP_CIPHER_CTX, decltype(&::EVP_CIPHER_CTX_free)> ctx(
        EVP_CIPHER_CTX_new(), ::EVP_CIPHER_CTX_free);
    if (ctx.get() == nullptr) {
        logWarning(10050) << "Failed to create EVP_CIPHER_CTX";
        return Streaming::ConstBuffer();
    }

    pool->reserve(clearText.size()); // result is the same size as the input
    const auto nonce = nonceFromTime(timestamp);
    assert(nonce.size() == 12);
    if (EVP_EncryptInit_ex(ctx.get(), EVP_chacha20(), nullptr, key.data(), nonce.data()) != 1) {
        return Streaming::ConstBuffer();
    }

    int len = 0;
    if (EVP_EncryptUpdate(ctx.get(),
            reinterpret_cast<uint8_t*>(pool->begin()),
            &len,
            reinterpret_cast<const uint8_t*>(clearText.begin()), clearText.size()) != 1) {
        logWarning(10050) << "Failed to create update encryption";
        return Streaming::ConstBuffer();
    }
    int ciphertext_len = len;
    if (EVP_EncryptFinal_ex(ctx.get(),
            reinterpret_cast<uint8_t*>(pool->begin() + len), &len) != 1) {
        logWarning(10050) << "Failed to finalize encryption";
        return Streaming::ConstBuffer();
    }
    ciphertext_len += len;
    return pool->commit(ciphertext_len);
}

Streaming::ConstBuffer decrypt(std::vector<uint8_t> &key, uint32_t timestamp, const Streaming::ConstBuffer &cipherText, const std::shared_ptr<Streaming::BufferPool> &pool)
{
    assert(key.size() == 32);

    std::unique_ptr<EVP_CIPHER_CTX, decltype(&::EVP_CIPHER_CTX_free)> ctx(
        EVP_CIPHER_CTX_new(), ::EVP_CIPHER_CTX_free);
    if (ctx.get() == nullptr) {
        logWarning(10050) << "Failed to create EVP_CIPHER_CTX";
        return Streaming::ConstBuffer();
    }

    const auto nonce = nonceFromTime(timestamp);
    assert(nonce.size() == 12);

    if (EVP_DecryptInit_ex(ctx.get(), EVP_chacha20(), nullptr, key.data(), nonce.data()) != 1) {
        logWarning(10050) << "Failed to initialize decryption";
        return Streaming::ConstBuffer();
    }

    int len = 0;
    if (EVP_DecryptUpdate(ctx.get(),
            reinterpret_cast<uint8_t*>(pool->begin()), &len,
            reinterpret_cast<const uint8_t*>(cipherText.begin()), cipherText.size()) != 1) {
        logWarning(10050) << "Failed to update decryption";
        return Streaming::ConstBuffer();
    }
    int plaintext_len = len;
    if (EVP_DecryptFinal_ex(ctx.get(),
            reinterpret_cast<uint8_t*>(pool->begin() + len), &len) != 1) {
        logWarning(10050) << "Failed to finalize decryption";
        return Streaming::ConstBuffer();
    }
    plaintext_len += len;
    return pool->commit(plaintext_len);
}

}


WalletData::WalletData(const std::shared_ptr<Wallet> &wallet, BackupSyncModuleInfo *parent)
    : QObject(parent),
    m_parent(parent),
    m_wallet(wallet)
{
    connect (m_wallet.get(), &Wallet::appendedTransactions, this, &WalletData::onTxAdded);
    connect (m_wallet.get(), &Wallet::transactionConfirmed, this, &WalletData::onTxConfirmed);
    connect (m_wallet.get(), &Wallet::transactionChanged, this, &WalletData::onTxUpdated);

    connect (this, SIGNAL(startUpload()), this, SLOT(upload()), Qt::QueuedConnection);
}

void WalletData::onTxUpdated(int txIndex)
{
    if (m_needsSave)
        return;
    auto info = m_wallet->fetchTransactionInfoSmall(txIndex);
    assert(info.exists());
    if (info.isUnconfirmed() || info.isRejected())
        return;
    // ensure we save user comments.
    // we skip mixTx's since the comment is auto-generated.
    if (!info.isMixTx && !info.userComment.isEmpty())
        markDirty();
}

void WalletData::onTxConfirmed(int txIndex)
{
    if (m_needsSave)
        return;
    auto info = m_wallet->fetchTransactionInfoSmall(txIndex);
    assert(info.exists());
    if (!info.isMixTx && !info.userComment.isEmpty())
        markDirty();
}

void WalletData::onTxAdded(int firstNew, int count)
{
    if (m_needsSave)
        return;
    for (int i = firstNew; i < firstNew + count; ++i) {
        auto info = m_wallet->fetchTransactionInfoSmall(i);
        assert(info.exists());
        if (info.isUnconfirmed() || info.isRejected())
            continue;
        if (!info.isMixTx && !info.userComment.isEmpty()) {
            markDirty();
            break;
        }
    }
}

WalletData::SaveData WalletData::save()
{
    /*
     * What we save is;
     *   a list of block heights where transactions are found relevant for our wallet.
     *   Interleaved the txid if needed, and then the transaction comment.
     *   Future: add metadata for p2sh or p2s style transactions, in order to be able
     *    to detect them on chain and thus later spend them.
     *    If it is already spent, maybe just a list of outputs that are 'mine' ?
     *   Last, a blockheight at which the wallet was saved. Implying that all newer
     *   blocks need checking.
     *
     * This is saved to a block which is then put in an envelope as an encrypted blob.
     * The envelope contains the timestamp, which is needed to decrypt the data.
     * The envelope is then also signed with the original identity resolved from the
     * wallets HD master key.
     */

    assert(m_wallet);
    assert(m_wallet->isHDWallet());
    SaveData saveData;
    if (!deriveAuthKey())
        return saveData;
    assert(m_privateKey.isValid());

    m_needsSave = false;

    int txIndex = 1;
    int blockHeight = 0;
    int bytes = 0;
    do {
        auto info = m_wallet->fetchTransactionInfoSmall(txIndex++);
        if (!info.exists())
            break;
        if (info.minedBlockHeight < 0)
            continue;
        if (info.minedBlockHeight != blockHeight) {
            bytes += 10;
            blockHeight = info.minedBlockHeight;
        }
        if (!info.userComment.isEmpty())
            bytes += 11 + info.userComment.size() * 3;
    } while (true);
    // TODO have a max size?

    // now actually do it.
    auto pool = Streaming::pool(bytes + 20000);
    Streaming::MessageBuilder builder(pool);
    txIndex = 1;
    do { // iterate over all transactions
        auto info = m_wallet->fetchTransactionInfoSmall(txIndex++);
        if (!info.exists())
            break;
        if (info.minedBlockHeight < 0)
            continue;
        if (info.minedBlockHeight != blockHeight) {
            blockHeight = info.minedBlockHeight;
            builder.add(BCHMDF::BlockHeight, blockHeight);
        }
        if (!info.userComment.isEmpty()) {
            builder.addByteArray(BCHMDF::TransactionIdShort, info.txid.begin(), 3);
            builder.add(BCHMDF::TransactionComment, info.userComment.toUtf8().constData());
        }
    } while (true);
    int lastSavedHeight = m_wallet->segment()->lastBlockSynched();
    builder.add(BCHMDF::LastBlockSeen, lastSavedHeight);

    auto theData = pool->commit();
    logDebug(10050) << "Data size:" << theData.size();
    if (theData.isEmpty())
        return saveData;

    CSHA256 messageHasher;
    messageHasher.write(theData.begin(), theData.size());
    char buf[CSHA256::OUTPUT_SIZE];
    messageHasher.finalize(buf);
    saveData.dataHash = uint256(buf);
    if (saveData.dataHash == m_lastSaveHash)
        return saveData;

    // we encrypt the data so only the owner can understand it.
    // step one, use the bitcoin private key (secp256k1) and derive a
    // private key from that to use for our symmetric crypto algo
    std::vector<uint8_t> derivedPrivKey = derivePrivKey(std::vector<uint8_t>(
        m_privateKey.begin(), m_privateKey.end()));
    if (derivedPrivKey.size() != 32) {
        logCritical(10050) << "Failed to derive priv key via HKDF";
        return saveData;
    }
    assert(derivedPrivKey.size() == 32);
    saveData.timestamp = QDateTime::currentSecsSinceEpoch();
    theData = encrypt(derivedPrivKey, saveData.timestamp, theData, pool);
    if (theData.isEmpty()) {
        logDebug(10050) << "encrypt failed";
        return saveData;
    }

    pool->reserve(100 + theData.size());
    pool->writeInt32(saveData.timestamp);
    pool->writeInt32(theData.size());
    pool->write(theData);

    // we need to create a base image to sign, the above data is what we'll use.
    CSHA256 hasher;
    hasher.write(pool->begin(), pool->size());
    hasher.finalize(buf);
    std::vector<unsigned char> signature;
    m_privateKey.signCompact(uint256(buf), signature);
    assert(signature.size() == 65);
    pool->write(signature.data(), signature.size());
    saveData.data = pool->commit();

#if 0
    QFile backup(m_wallet->name() + "_backup");
    if (backup.open(QIODevice::WriteOnly))
        backup.write(saveData.data.begin(), saveData.data.size());
    else
        logCritical(10050) << "Failed to open file" << backup.fileName();
#endif

    return saveData;
}

void WalletData::upload()
{
    if (m_httpClient.get())
        return;
    if (FloweePay::instance()->isOffline())
        return;
    AccountConfig accountConfig(m_wallet); // honor user config
    if (!accountConfig.cloudStoreDetails())
        return;

    auto saveData = save();
    if (saveData.data.isEmpty())
        return;


    // Prepare the HTTP request
    boost::beast::http::request<boost::beast::http::vector_body<char>> request;
    request.version(11);
    request.method(boost::beast::http::verb::post);
    request.target("/md");
    request.set(boost::beast::http::field::host, "flowee.org");
    request.set(boost::beast::http::field::content_type, "application/octet-stream");
    request.body() = std::vector<char>(saveData.data.begin(), saveData.data.end());
    auto *fp = FloweePay::instance();
    m_httpClient = SimpleHttpClient::create(fp->ioContext(), fp->sslContext(), request);
    connect (m_httpClient.get(), &SimpleHttpClient::finished, this, [=]() {
        if (m_httpClient->error() == SimpleHttpClient::NoError) {
            m_lastSaveHash = saveData.dataHash;
            setLastSavedTime(saveData.timestamp); // remember 'when' we saved.
            emit m_parent->doRequestSave();
        }
        m_httpClient.reset();
    });
}

void WalletData::restore(RestoreReason reason)
{
    auto start = m_wallet->walletCreatedHeight();
    if (start < 0) { // negative numbers mean it is unused.
        if (reason == BeforeSave)
            emit startUpload();
        return;
    }
    auto now = FloweePay::instance()->headerChainHeight();
    if (now - start < 2 && m_wallet->lastTransactionTimestamp() == 0) {
        // wallet is not an imported seed, and is empty. Ignore
        if (reason == BeforeSave)
            emit startUpload();
        return;
    }
    if (!deriveAuthKey()) {
        // basically only fails on an encrypted wallet
        return;
    }
    // if wallet is not at 'tip', return later
    const int headerChainHeight = FloweePay::instance()->headerChainHeight();
    if (m_wallet->segment()->lastBlockSynched() < headerChainHeight) {
        emit storeWaitingForSync();
        logCritical(10050) << "Delaying restore/upload wallet: not yet at tip.";
        QTimer::singleShot(10 * 1000, this, [=]() { restore(reason); } );
        return;
    }
    auto authAddress = fancyAddess();

    // check online
    boost::beast::http::request<boost::beast::http::vector_body<char>> request;
    request.version(11);
    request.method(boost::beast::http::verb::get);
    request.target(QString("/md/%1").arg(fancyAddess()).toStdString());
    request.set(boost::beast::http::field::host, "flowee.org");
    auto *fp = FloweePay::instance();
    m_httpClient = SimpleHttpClient::create(fp->ioContext(), fp->sslContext(), request);
    connect (m_httpClient.get(), &SimpleHttpClient::finished, this, [=]() {
        auto copySharedPtr = m_httpClient;
        m_httpClient.reset();
        if (copySharedPtr->error() == SimpleHttpClient::NoError) {
            auto buf = copySharedPtr->responseBody();
            const auto doc = QJsonDocument::fromJson(QByteArray(buf.begin(), buf.size()));
            const auto root = doc.object();
            const auto list = root.keys();
            uint64_t latest = 0;
            QString latestStr;
            for (const auto &item : list) {
                auto timestampStr = root[item].toString();
                bool ok;
                uint64_t timestamp = timestampStr.toLong(&ok);
                if (ok && timestamp > latest) {
                    latest = timestamp;
                    latestStr = timestampStr;
                }
            }
            // check to avoid downloading, if latest one was generated by us.
            if (latest != m_lastSavedTime)
                restoreFrom(latestStr);
        }
        // Notice that if the server has never heard of this wallet before, it returns
        // http for (or SimpleHttpClient::PeerIssues). As such errors are expected and Ok.
        // We schedule an upload if that was what the user wanted now.
        if (reason == BeforeSave)
            emit startUpload();
    });
}

void WalletData::markDirty()
{
    if (m_needsSave)
        return;
    m_needsSave = true;
    // we'll do an upload after restore finished.
    // see connect in constructor.
    QTimer::singleShot(120 * 1000, this, [=]() { restore(BeforeSave); });
    emit m_parent->doRequestSave();
}

bool WalletData::deriveAuthKey()
{
    if (!m_privateKey.isValid()) {
        if (!m_wallet->isDecrypted()) {
            logInfo(10050) << "Skipping encrypted wallet" << m_wallet->name();
            return false;
        }
        HDMasterKey mk = HDMasterKey::fromMnemonic(m_wallet->hdWalletMnemonic().toStdString(),
                                                   HDMasterKey::BIP39Mnemonic);

        auto pathVector = HDMasterKey::deriveFromString(m_wallet->derivationPath().toStdString());
        assert(pathVector.size() == 3);
        pathVector.resize(5);
        pathVector[3] = 3; // the Ownership chain.
        pathVector[4] = 0; // First key is reserved for backups.
        m_privateKey = mk.derive(pathVector);
    }
    return true;
}

void WalletData::restoreFrom(const QString &timestamp)
{
    if (m_httpClient.get())
        return;
    boost::beast::http::request<boost::beast::http::vector_body<char>> request;
    request.version(11);
    request.method(boost::beast::http::verb::get);
    request.target(QString("/md/%1/%2").arg(fancyAddess(), timestamp).toStdString());
    request.set(boost::beast::http::field::host, "flowee.org");
    auto *fp = FloweePay::instance();
    m_httpClient = SimpleHttpClient::create(fp->ioContext(), fp->sslContext(), request);
    connect (m_httpClient.get(), &SimpleHttpClient::finished, this, [=]() {
        if (m_httpClient->error() == SimpleHttpClient::NoError)
            restoreDownload(m_httpClient->responseBody());
        m_httpClient.reset();
    });
}

void WalletData::restoreDownload(const Streaming::ConstBuffer &blob)
{
    if (blob.size() < 77) {
        logCritical(10050) << "The downloaded backup is too small." << blob.size();
        return;
    }

    auto timestampData = blob.mid(0, 4);
    uint32_t timestamp = *reinterpret_cast<const uint32_t*>(timestampData.begin());
    auto dataSizeData = blob.mid(4, 4);
    uint32_t dataSize = *reinterpret_cast<const uint32_t*>(dataSizeData.begin());
    if (blob.size() - 65 - 4 - 4 != (int) dataSize) {
        logCritical(10050) << "Data-size invalid";
        return;
    }
    auto data = blob.mid(8, dataSize);
    auto signature = blob.mid(8 + dataSize, 65);

    PublicKey pk;
    CSHA256 hasher;
    hasher.write(blob.begin(), 8 + dataSize);
    char buf[CSHA256::OUTPUT_SIZE];
    hasher.finalize(buf);
    uint256 hash(buf);
    std::vector<uint8_t> sig(signature.begin(), signature.end());
    if (!pk.recoverCompact(hash, sig)) {
        logWarning(10050) << "Download invalid";
        return;
    }
    deriveAuthKey();
    if (pk != m_privateKey.getPubKey()) {
        logWarning(10050) << "Download not created by us";
        return;
    }
    if (!pk.verifyCompact(hash, sig)) {
        logWarning(10050) << "Downloads' signature invalid";
        return;
    }

    std::vector<uint8_t> derivedPrivKey = derivePrivKey(std::vector<uint8_t>(
        m_privateKey.begin(), m_privateKey.end()));
    auto message = decrypt(derivedPrivKey, timestamp, data, Streaming::pool(0));
    Streaming::MessageParser parser(message);
    // Today we simply apply comments, nothing more.
    Streaming::ConstBuffer txIdShort;
    int txIndex = 1;
    int txIndexInBlock = 1;
    int height = -1;
    while (parser.next() == Streaming::FoundTag) {
        if (parser.tag() == BCHMDF::BlockHeight) {
            txIdShort.clear();
            height = parser.intData();
            txIndexInBlock = txIndex;
        }
        else if (parser.tag() == BCHMDF::TransactionIdShort) {
            txIdShort = parser.bytesDataBuffer();
        }
        else if (parser.tag() == BCHMDF::TransactionComment) {
            txIndex = txIndexInBlock; // check all transactions in block for match.
            while (txIdShort.size() > 1 && txIdShort.size() <= 32) { // 32 is sizeof uint256
                auto info = m_wallet->fetchTransactionInfoSmall(txIndex);
                if (!info.exists())
                    break;
                assert(txIdShort.size() <= (int) info.txid.size());
                if (info.minedBlockHeight > height)
                    break;
                if (info.minedBlockHeight < height)
                    txIndexInBlock = txIndex;
                if (info.minedBlockHeight == height
                        && memcmp(info.txid.begin(), txIdShort.begin(), txIdShort.size()) == 0) {
                    // match!
                    if (info.userComment.isEmpty()) { // don't overwrite
                        auto bytes = parser.bytesDataBuffer();
                        m_wallet->setTransactionComment(txIndex,
                                QString::fromUtf8(bytes.begin(), bytes.size()));
                    }
                    ++txIndex;
                    break;
                }
                ++txIndex;
            }
        }
    }
}

uint32_t WalletData::lastSavedTime() const
{
    return m_lastSavedTime;
}

void WalletData::setLastSavedTime(uint32_t newLastSavedTime)
{
    m_lastSavedTime = newLastSavedTime;
    emit lastSavedChanged();
}

bool WalletData::needsSave() const
{
    return m_needsSave;
}

std::shared_ptr<Wallet> WalletData::wallet() const
{
    return m_wallet;
}

void WalletData::setPrivKeyData(const Streaming::ConstBuffer &data)
{
    m_privateKey.set(
        reinterpret_cast<const uint8_t*>(data.begin()),
        reinterpret_cast<const uint8_t*>(data.end()));
}

const PrivateKey &WalletData::privKey() const
{
    return m_privateKey;
}

QString WalletData::fancyAddess() const
{
    if (!m_privateKey.isValid())
        return QString();
    auto id = m_privateKey.getPubKey().getKeyId();
    return renderAddress(id);
}