// Copyright (c) 2012-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 <bench/bench.h>
#include <consensus/amount.h>
#include <interfaces/chain.h>
#include <node/context.h>
#include <outputtype.h>
#include <policy/feerate.h>
#include <policy/policy.h>
#include <primitives/transaction.h>
#include <random.h>
#include <sync.h>
#include <util/result.h>
#include <wallet/coinselection.h>
#include <wallet/spend.h>
#include <wallet/test/util.h>
#include <wallet/transaction.h>
#include <wallet/wallet.h>

#include <cassert>
#include <map>
#include <memory>
#include <set>
#include <utility>
#include <vector>
#include <cmath>

using node::NodeContext;
using wallet::AttemptSelection;
using wallet::CHANGE_LOWER;
using wallet::COutput;
using wallet::CWallet;
using wallet::CWalletTx;
using wallet::CoinEligibilityFilter;
using wallet::CoinSelectionParams;
using wallet::CreateMockableWalletDatabase;
using wallet::OutputGroup;
using wallet::SelectCoinsBnB;
using wallet::TxStateInactive;

namespace {
constexpr CAmount operator""_sat(uint64_t v) { return v; }
constexpr CAmount operator""_ksat(uint64_t v) { return v * 1'000; }
constexpr CAmount operator""_btc(uint64_t v) { return v * COIN; }
constexpr CAmount operator""_btc(long double v) { return std::llround(v * COIN); }

int rand_percentage(FastRandomContext& rng) { return rng.randrange(100); }
CAmount rand_range(FastRandomContext& rng, CAmount lo, CAmount hi) { return lo + rng.randrange(hi - lo + 1); }
} // namespace

static void addCoin(CAmount nValue, std::vector<std::unique_ptr<CWalletTx>>& wtxs)
{
    static int nextLockTime = 0;
    CMutableTransaction tx;
    tx.nLockTime = nextLockTime++; // so all transactions get different hashes
    tx.vout.resize(1);
    tx.vout[0].nValue = nValue;
    wtxs.push_back(std::make_unique<CWalletTx>(MakeTransactionRef(std::move(tx)), TxStateInactive{}));
}

// Simple benchmark for wallet coin selection that exercises a worst-case
// scenario for Knapsack: All UTXOs are necessary, but it is not an exact
// match, so the only eligible input set is only discovered on the second pass
// after all random walks fail to produce a solution.
static void KnapsackWorstCase(benchmark::Bench& bench)
{
    NodeContext node;
    auto chain{interfaces::MakeChain(node)};
    CWallet wallet(chain.get(), "", CreateMockableWalletDatabase());
    std::vector<std::unique_ptr<CWalletTx>> wtxs;
    LOCK(wallet.cs_wallet);

    for (int i{0}; i < 1000; ++i) {
        addCoin(1000_btc, wtxs);
    }
    addCoin(3_btc, wtxs);

    // Create coins
    wallet::CoinsResult available_coins;
    for (const auto& wtx : wtxs) {
        const auto txout{wtx->tx->vout.at(0)};
        const int input_bytes{CalculateMaximumSignedInputSize(txout, &wallet, /*coin_control=*/nullptr)};
        const COutput output{COutPoint(wtx->GetHash(), 0), txout, /*depth=*/6 * 24, input_bytes, /*solvable=*/true, /*safe=*/true, wtx->GetTxTime(), /*from_me=*/true, /*fees=*/0};
        available_coins.coins[OutputType::BECH32].emplace_back(output);
    }

    const CoinEligibilityFilter filter_standard(1, 6, 0);
    FastRandomContext rand{/*fDeterministic=*/true};
    const CoinSelectionParams coin_selection_params{
        rand,
        /*change_output_size=*/34,
        /*change_spend_size=*/148,
        /*min_change_target=*/CHANGE_LOWER,
        /*effective_feerate=*/CFeeRate(20'000),
        /*long_term_feerate=*/CFeeRate(10'000),
        /*discard_feerate=*/CFeeRate(3000),
        /*tx_noinputs_size=*/0,
        /*avoid_partial=*/false,
    };
    auto group{wallet::GroupOutputs(wallet, available_coins, coin_selection_params, {{filter_standard}})[filter_standard]};
    bench.run([&] {
        auto result{AttemptSelection(wallet.chain(), 1002.99_btc, group, coin_selection_params, /*allow_mixed_output_types=*/true)};
        assert(result);
        assert(result->GetSelectedValue() == 1003_btc);
        assert(result->GetInputSet().size() == 2);
    });
}

// This benchmark is based on a UTXO pool composed of 400 UTXOs. The UTXOs are
// pseudorandomly generated to be of the four relevant output types P2PKH,
// P2SH-P2WPKH, P2WPKH, and P2TR UTXOs, and fall in the range of 10'000 sats to
// 1₿ with larger amounts being more likely.
// This UTXO pool is used to run coin selection for 100 pseudorandom selection
// targets from 0.1–2₿. Altogether, this gives us a deterministic benchmark
// with a hopefully somewhat representative coin selection scenario.
static void CoinSelectionOnDiverseWallet(benchmark::Bench& bench)
{
    FastRandomContext rng{/*fDeterministic=*/true};

    NodeContext node;
    auto chain{interfaces::MakeChain(node)};
    CWallet wallet(chain.get(), "", CreateMockableWalletDatabase());
    LOCK(wallet.cs_wallet);

    std::vector<std::unique_ptr<CWalletTx>> wtxs;
    wtxs.reserve(400);
    for (size_t i{0}; i < wtxs.capacity(); ++i) {
        const int p{rand_percentage(rng)};
        const auto val{p < 50 ? rand_range(rng, 10_ksat, 100_ksat) :
                       p < 75 ? rand_range(rng, 100_ksat, 1000_ksat) :
                       p < 95 ? rand_range(rng, 1000_ksat, 1_btc) :
                                rand_range(rng, 0.1_btc, 1_btc)};
        addCoin(val, wtxs);
    }

    // Create coins
    wallet::CoinsResult available_coins;
    for (const auto& wtx : wtxs) {
        const int p{rand_percentage(rng)};
        auto val{p < 35 ? OutputType::LEGACY :
                 p < 55 ? OutputType::P2SH_SEGWIT :
                 p < 90 ? OutputType::BECH32 :
                          OutputType::BECH32M};

        const auto txout{wtx->tx->vout.at(0)};
        const int input_bytes{CalculateMaximumSignedInputSize(txout, &wallet, /*coin_control=*/nullptr)};
        const COutput output{COutPoint{wtx->GetHash(), 0}, txout, /*depth=*/6 * 24, input_bytes, /*solvable=*/true, /*safe=*/true, wtx->GetTxTime(), /*from_me=*/true, /*fees=*/0};
        available_coins.coins[val].emplace_back(output);
    }

    std::vector<CAmount> targets;
    targets.reserve(10);
    for (size_t i{0}; i < targets.capacity(); ++i) {
        targets.push_back(rand_range(rng, 0.1_btc, 1_btc));
    }

    const CoinSelectionParams coin_selection_params{
        rng,
        /*change_output_size=*/31,
        /*change_spend_size=*/68,
        /*min_change_target=*/CHANGE_LOWER,
        /*effective_feerate=*/CFeeRate(20'000),
        /*long_term_feerate=*/CFeeRate(10'000),
        /*discard_feerate=*/CFeeRate(3000),
        /*tx_noinputs_size=*/72,
        /*avoid_partial=*/false,
    };
    const CoinEligibilityFilter filter_standard(1, 6, 0);
    auto group{wallet::GroupOutputs(wallet, available_coins, coin_selection_params, {{filter_standard}})[filter_standard]};

    bench.run([&] {
        for (const auto& target : targets) {
            auto result{AttemptSelection(wallet.chain(), target, group, coin_selection_params, /*allow_mixed_output_types=*/true)};
            assert(result && result->GetSelectedValue() >= target);
        }
    });
}

static void add_coin(CAmount nValue, uint32_t nInput, std::vector<OutputGroup>& set)
{
    CMutableTransaction tx;
    tx.vout.resize(nInput + 1);
    tx.vout[nInput].nValue = nValue;
    COutput output(COutPoint{tx.GetHash(), nInput}, tx.vout.at(nInput), /*depth=*/0, /*input_bytes=*/-1, /*solvable=*/true, /*safe=*/true, /*time=*/0, /*from_me=*/true, /*fees=*/0);
    set.emplace_back();
    set.back().Insert(std::make_shared<COutput>(output), /*ancestors=*/0, /*descendants=*/0);
}

static CAmount make_hard_case(int utxos, std::vector<OutputGroup>& utxo_pool)
{
    utxo_pool.clear();
    CAmount target{0};
    for (int i{0}; i < utxos; ++i) {
        target += 1_sat << (utxos + i);
        add_coin(1_sat << (utxos + i), 2 * i, utxo_pool);
        add_coin((1_sat << (utxos + i)) + (1_sat << (utxos - 1 - i)), 2 * i + 1, utxo_pool);
    }
    return target;
}

static void BnBExhaustion(benchmark::Bench& bench)
{
    std::vector<OutputGroup> utxo_pool;
    auto target{make_hard_case(17, utxo_pool)};
    bench.run([&] {
        SelectCoinsBnB(utxo_pool, target, 0, MAX_STANDARD_TX_WEIGHT); // Should exhaust
    });
}

BENCHMARK(KnapsackWorstCase, benchmark::PriorityLevel::HIGH);
BENCHMARK(CoinSelectionOnDiverseWallet, benchmark::PriorityLevel::HIGH);
BENCHMARK(BnBExhaustion, benchmark::PriorityLevel::HIGH);