// Copyright (c) 2022 The Bitcoin Core developers
// Distributed under the MIT software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.

#include <kernel/mempool_persist.h>

#include <clientversion.h>
#include <consensus/amount.h>
#include <logging.h>
#include <primitives/transaction.h>
#include <random.h>
#include <serialize.h>
#include <streams.h>
#include <sync.h>
#include <txmempool.h>
#include <uint256.h>
#include <util/fs.h>
#include <util/fs_helpers.h>
#include <util/signalinterrupt.h>
#include <util/time.h>
#include <validation.h>

#include <cstdint>
#include <cstdio>
#include <exception>
#include <functional>
#include <map>
#include <memory>
#include <set>
#include <stdexcept>
#include <utility>
#include <vector>

using fsbridge::FopenFn;

namespace kernel {

static const uint64_t MEMPOOL_DUMP_VERSION_NO_XOR_KEY{1};
static const uint64_t MEMPOOL_DUMP_VERSION{2};

bool LoadMempool(CTxMemPool& pool, const fs::path& load_path, Chainstate& active_chainstate, ImportMempoolOptions&& opts)
{
    if (load_path.empty()) return false;

    AutoFile file{opts.mockable_fopen_function(load_path, "rb")};
    if (file.IsNull()) {
        LogPrintf("Failed to open mempool file from disk. Continuing anyway.\n");
        return false;
    }

    int64_t count = 0;
    int64_t expired = 0;
    int64_t failed = 0;
    int64_t already_there = 0;
    int64_t unbroadcast = 0;
    const auto now{NodeClock::now()};

    try {
        uint64_t version;
        file >> version;
        std::vector<std::byte> xor_key;
        if (version == MEMPOOL_DUMP_VERSION_NO_XOR_KEY) {
            // Leave XOR-key empty
        } else if (version == MEMPOOL_DUMP_VERSION) {
            file >> xor_key;
        } else {
            return false;
        }
        file.SetXor(xor_key);
        uint64_t num;
        file >> num;
        while (num) {
            --num;
            CTransactionRef tx;
            int64_t nTime;
            int64_t nFeeDelta;
            file >> TX_WITH_WITNESS(tx);
            file >> nTime;
            file >> nFeeDelta;

            if (opts.use_current_time) {
                nTime = TicksSinceEpoch<std::chrono::seconds>(now);
            }

            CAmount amountdelta = nFeeDelta;
            if (amountdelta && opts.apply_fee_delta_priority) {
                pool.PrioritiseTransaction(tx->GetHash(), amountdelta);
            }
            if (nTime > TicksSinceEpoch<std::chrono::seconds>(now - pool.m_expiry)) {
                LOCK(cs_main);
                const auto& accepted = AcceptToMemoryPool(active_chainstate, tx, nTime, /*bypass_limits=*/false, /*test_accept=*/false);
                if (accepted.m_result_type == MempoolAcceptResult::ResultType::VALID) {
                    ++count;
                } else {
                    // mempool may contain the transaction already, e.g. from
                    // wallet(s) having loaded it while we were processing
                    // mempool transactions; consider these as valid, instead of
                    // failed, but mark them as 'already there'
                    if (pool.exists(GenTxid::Txid(tx->GetHash()))) {
                        ++already_there;
                    } else {
                        ++failed;
                    }
                }
            } else {
                ++expired;
            }
            if (active_chainstate.m_chainman.m_interrupt)
                return false;
        }
        std::map<uint256, CAmount> mapDeltas;
        file >> mapDeltas;

        if (opts.apply_fee_delta_priority) {
            for (const auto& i : mapDeltas) {
                pool.PrioritiseTransaction(i.first, i.second);
            }
        }

        std::set<uint256> unbroadcast_txids;
        file >> unbroadcast_txids;
        if (opts.apply_unbroadcast_set) {
            unbroadcast = unbroadcast_txids.size();
            for (const auto& txid : unbroadcast_txids) {
                // Ensure transactions were accepted to mempool then add to
                // unbroadcast set.
                if (pool.get(txid) != nullptr) pool.AddUnbroadcastTx(txid);
            }
        }
    } catch (const std::exception& e) {
        LogPrintf("Failed to deserialize mempool data on disk: %s. Continuing anyway.\n", e.what());
        return false;
    }

    LogPrintf("Imported mempool transactions from disk: %i succeeded, %i failed, %i expired, %i already there, %i waiting for initial broadcast\n", count, failed, expired, already_there, unbroadcast);
    return true;
}

bool DumpMempool(const CTxMemPool& pool, const fs::path& dump_path, FopenFn mockable_fopen_function, bool skip_file_commit)
{
    auto start = SteadyClock::now();

    std::map<uint256, CAmount> mapDeltas;
    std::vector<TxMempoolInfo> vinfo;
    std::set<uint256> unbroadcast_txids;

    static Mutex dump_mutex;
    LOCK(dump_mutex);

    {
        LOCK(pool.cs);
        for (const auto &i : pool.mapDeltas) {
            mapDeltas[i.first] = i.second;
        }
        vinfo = pool.infoAll();
        unbroadcast_txids = pool.GetUnbroadcastTxs();
    }

    auto mid = SteadyClock::now();

    AutoFile file{mockable_fopen_function(dump_path + ".new", "wb")};
    if (file.IsNull()) {
        return false;
    }

    try {
        const uint64_t version{pool.m_persist_v1_dat ? MEMPOOL_DUMP_VERSION_NO_XOR_KEY : MEMPOOL_DUMP_VERSION};
        file << version;

        std::vector<std::byte> xor_key(8);
        if (!pool.m_persist_v1_dat) {
            FastRandomContext{}.fillrand(xor_key);
            file << xor_key;
        }
        file.SetXor(xor_key);

        file << (uint64_t)vinfo.size();
        for (const auto& i : vinfo) {
            file << TX_WITH_WITNESS(*(i.tx));
            file << int64_t{count_seconds(i.m_time)};
            file << int64_t{i.nFeeDelta};
            mapDeltas.erase(i.tx->GetHash());
        }

        file << mapDeltas;

        LogPrintf("Writing %d unbroadcast transactions to disk.\n", unbroadcast_txids.size());
        file << unbroadcast_txids;

        if (!skip_file_commit && !FileCommit(file.Get()))
            throw std::runtime_error("FileCommit failed");
        file.fclose();
        if (!RenameOver(dump_path + ".new", dump_path)) {
            throw std::runtime_error("Rename failed");
        }
        auto last = SteadyClock::now();

        LogPrintf("Dumped mempool: %gs to copy, %gs to dump\n",
                  Ticks<SecondsDouble>(mid - start),
                  Ticks<SecondsDouble>(last - mid));
    } catch (const std::exception& e) {
        LogPrintf("Failed to dump mempool: %s. Continuing anyway.\n", e.what());
        return false;
    }
    return true;
}

} // namespace kernel