// Copyright (c) 2016-2024 The Hush developers
// Copyright (c) 2019 CryptoForge
// Distributed under the GPLv3 software license, see the accompanying
// file COPYING or https://www.gnu.org/licenses/gpl-3.0.en.html

#include "assert.h"
#include "boost/variant/static_visitor.hpp"
#include "asyncrpcoperation_saplingconsolidation.h"
#include "init.h"
#include "key_io.h"
#include "rpc/protocol.h"
#include "random.h"
#include "sync.h"
#include "tinyformat.h"
#include "transaction_builder.h"
#include "util.h"
#include "utilmoneystr.h"
#include "wallet.h"

CAmount fConsolidationTxFee = DEFAULT_CONSOLIDATION_FEE;
bool fConsolidationMapUsed = false;

// Number of Sietch dummy ("zdust") shielded outputs added to every consolidation
// transaction to obscure the real output and keep the anonymity set large. This is
// a wallet privacy-tuning parameter (not a consensus rule); the sweep operation uses
// the same value under the name ZOUTS.
static const int MIN_ZOUTS = 7;

extern string randomSietchZaddr();

AsyncRPCOperation_saplingconsolidation::AsyncRPCOperation_saplingconsolidation(int targetHeight) : targetHeight_(targetHeight) {}

AsyncRPCOperation_saplingconsolidation::~AsyncRPCOperation_saplingconsolidation() {}

void AsyncRPCOperation_saplingconsolidation::main() {
    if (isCancelled()) {
        // Only the current op owns the scheduler flag; a stale/cancelled op must
        // not clear it out from under a freshly-enqueued successor.
        if (pwalletMain) {
            LOCK(pwalletMain->cs_wallet);
            if (getId() == pwalletMain->saplingConsolidationOperationId) {
                pwalletMain->fConsolidationRunning = false;
            }
        }
        return;
    }

    set_state(OperationStatus::EXECUTING);
    start_execution_clock();

    bool success = false;

    try {
        success = main_impl();
    } catch (...) {
        set_error_from_current_exception();
    }

    stop_execution_clock();

    if (success) {
        set_state(OperationStatus::SUCCESS);
    } else {
        set_state(OperationStatus::FAILED);
    }

    std::string s = strprintf("%s: Sapling Consolidation operation complete. (status=%s", getId(), getStateAsString());
    if (success) {
        s += strprintf(", success)\n");
    } else {
        s += strprintf(", error=%s)\n", getErrorMessage());
    }

    LogPrintf("%s", s);
    unlock_notes(); // clean up
    LogPrint("zrpc", "%s: consolidation input notes unlocked\n", getId());

    // Advance the interval and clear the running flag on EVERY terminal state
    // (success, failure, exception) so consolidation runs once per interval
    // instead of every block, and a failed round still lets the next one fire.
    // Only the CURRENT op does this bookkeeping. This fixes the pre-existing
    // wedge where nextConsolidation never advanced and fConsolidationRunning
    // was never set/reset.
    if (pwalletMain) {
        LOCK2(cs_main, pwalletMain->cs_wallet);
        if (getId() == pwalletMain->saplingConsolidationOperationId) {
            int tipHeight = (chainActive.Tip() != NULL) ? chainActive.Tip()->GetHeight() : targetHeight_;
            pwalletMain->nextConsolidation     = pwalletMain->consolidationInterval + tipHeight;
            pwalletMain->fConsolidationRunning = false;
        }
    }
}

bool AsyncRPCOperation_saplingconsolidation::main_impl() {
    bool status=true;
    auto opid=getId();
    LogPrintf("%s: Beginning AsyncRPCOperation_saplingconsolidation\n", opid);
    auto consensusParams = Params().GetConsensus();
    int tipHeight;
    {
        LOCK(cs_main);
        tipHeight = (chainActive.Tip() != NULL) ? chainActive.Tip()->GetHeight() : targetHeight_;
    }

    // Build and expire against tipHeight (execution-time), not the stale
    // enqueue-time targetHeight_, so the builder's consensus-branch selection and
    // the NU-straddle guard agree with the height the tx is signed for. Mirrors
    // the autoshield op (commit 65130c312).
    auto nextActivationHeight = NextActivationHeight(tipHeight, consensusParams);
    if (nextActivationHeight && tipHeight + AUTO_OP_EXPIRY_DELTA >= nextActivationHeight.get()) {
        LogPrintf("%s: Consolidation txs would be created before a NU activation but may expire after. Skipping this round.\n",opid);
        setConsolidationResult(0, 0, std::vector<std::string>());
        return status;
    }

    std::vector<SaplingNoteEntry> saplingEntries;
    std::set<libzcash::SaplingPaymentAddress> addresses;
    {
        LOCK2(cs_main, pwalletMain->cs_wallet);
        // We set minDepth to 11 to avoid unconfirmed notes and in anticipation of specifying
        // an anchor at height N-10 for each SpendDescription
        // Consider, should notes be sorted?
        pwalletMain->GetFilteredNotes(saplingEntries, "", 11);

        if(saplingEntries.size() == 0) {
            LogPrintf("%s: Nothing to consolidate, done.\n",opid);
            return true;
        }

        // store sapling inputs so we can correctly lock and unlock them
        for (auto entry : saplingEntries) {
            z_sapling_inputs_.push_back(entry);
        }

        if (fConsolidationMapUsed) {
            const vector<string>& v = mapMultiArgs["-consolidatesaplingaddress"];
            for(int i = 0; i < v.size(); i++) {
                auto zAddress = DecodePaymentAddress(v[i]);
                if (boost::get<libzcash::SaplingPaymentAddress>(&zAddress) != nullptr) {
                    libzcash::SaplingPaymentAddress saplingAddress = boost::get<libzcash::SaplingPaymentAddress>(zAddress);
                    addresses.insert(saplingAddress);
                } else {
                    LogPrintf("%s: Invalid zaddr, exiting\n", opid);
                    return false;
                }
            }
        } else {
            pwalletMain->GetSaplingPaymentAddresses(addresses);
        }
    }

    int numTxCreated = 0;
    std::vector<std::string> consolidationTxIds;
    CAmount amountConsolidated = 0;
    CCoinsViewCache coinsView(pcoinsTip);

    for (auto addr : addresses) {
        libzcash::SaplingExtendedSpendingKey extsk;
        if (pwalletMain->GetSaplingExtendedSpendingKey(addr, extsk)) {

            std::vector<SaplingNoteEntry> fromNotes;
            CAmount amountToSend = 0;
            // max of 8 zins means the tx cannot reduce the anonset,
            // since there will be 8 zins and 8 zouts at worst case
            // This also helps reduce ztx creation time
            int maxQuantity = rand() % 8 + 1;
            for (const SaplingNoteEntry& saplingEntry : saplingEntries) {

                libzcash::SaplingIncomingViewingKey ivk;
                pwalletMain->GetSaplingIncomingViewingKey(boost::get<libzcash::SaplingPaymentAddress>(saplingEntry.address), ivk);

                //Select Notes from that same address we will be sending to.
                if (ivk == extsk.expsk.full_viewing_key().in_viewing_key()) {
                  amountToSend += CAmount(saplingEntry.note.value());
                  fromNotes.push_back(saplingEntry);
                }

                //Only use a randomly determined number of notes
                if (fromNotes.size() >= maxQuantity)
                  break;

            }

            // minimum required
            // We use 3 so that addresses can spent one zutxo and still have another zutxo to use while that
            // tx is confirming
            int minQuantity = 3;
            if (fromNotes.size() < minQuantity)
              continue;

            auto builder = TransactionBuilder(consensusParams, tipHeight, pwalletMain);
            builder.SetExpiryHeight(tipHeight + AUTO_OP_EXPIRY_DELTA);
            auto actualAmountToSend = amountToSend < fConsolidationTxFee ? 0 : amountToSend - fConsolidationTxFee;
            LogPrintf("%s: %s Beginning to create transaction with Sapling output amount=%s\n", __func__, opid, FormatMoney(actualAmountToSend));

            // Select Sapling notes
            std::vector<SaplingOutPoint> ops;
            std::vector<libzcash::SaplingNote> notes;
            for (auto fromNote : fromNotes) {
                ops.push_back(fromNote.op);
                notes.push_back(fromNote.note);
            }

            // Fetch Sapling anchor and witnesses
            uint256 anchor;
            std::vector<boost::optional<SaplingWitness>> witnesses;
            {
                LOCK2(cs_main, pwalletMain->cs_wallet);
                // LogPrint("zrpcunsafe", "%s: Fetching note witnesses\n", opid);
                pwalletMain->GetSaplingNoteWitnesses(ops, witnesses, anchor);
            }

            // Add Sapling spends
            for (size_t i = 0; i < notes.size(); i++) {
                if (!witnesses[i]) {
                    LogPrintf("%s: Missing Witnesses! Stopping.\n", opid);
                    status=false;
                    break;
                }
                builder.AddSaplingSpend(extsk.expsk, notes[i], anchor, witnesses[i].get());
                LogPrint("zrpcunsafe", "%s: Added consolidation input %d\n", opid, i);
            }

            CAmount thisTxFee = amountToSend < fConsolidationTxFee ? 0 : fConsolidationTxFee;
            LogPrint("zrpcunsafe", "%s: Using fee=%d\n", opid, thisTxFee);
            builder.SetFee(thisTxFee);

            // Add the actual consolidation tx
            builder.AddSaplingOutput(extsk.expsk.ovk, addr, actualAmountToSend);
            LogPrint("zrpcunsafe", "%s: Added consolidation output %s with amount=%li\n", opid, addr.GetHash().ToString().c_str(), actualAmountToSend);

            // Add sietch zouts: MIN_ZOUTS dummy zero-value shielded outputs to
            // randomly-generated z-addresses, so the consolidation tx does not
            // shrink the anonymity set.
            for(size_t i = 0; i < MIN_ZOUTS; i++) {
                string zdust = randomSietchZaddr();
                auto zaddr   = DecodePaymentAddress(zdust);
                if (IsValidPaymentAddress(zaddr)) {
                    CAmount amount=0;
                    auto sietchZoutput = boost::get<libzcash::SaplingPaymentAddress>(zaddr);
                    LogPrint("zrpcunsafe", "%s: Adding Sietch zdust output %d %s amount=%li\n", opid, i, zdust, amount);

                    // actually add our sietch zoutput, the new way
                    builder.AddSaplingOutput(extsk.expsk.ovk, sietchZoutput, amount);
                } else {
                    LogPrintf("%s: Invalid payment address %s! Stopping.\n", opid, zdust);
                    status = false;
                    break;
                }
            }
            LogPrint("zrpcunsafe", "%s: Done adding %d sietch zouts\n", opid, MIN_ZOUTS);
            auto maybe_tx = builder.Build();
            if (!maybe_tx) {
                LogPrintf("%s: Failed to build transaction.\n",opid);
                status=false;
                break;
            }
            CTransaction tx = maybe_tx.get();

            if (isCancelled()) {
                LogPrintf("%s: Canceled. Stopping.\n", opid);
                status=false;
                break;
            }

            // Lock shielded input notes
            lock_notes();
            LogPrint("zrpc", "%s: consolidation input notes locked\n", getId());

            if(pwalletMain->CommitAutomatedTx(tx)) {
                LogPrintf("%s: Committed consolidation transaction with txid=%s\n",opid, tx.GetHash().ToString());
                amountConsolidated += actualAmountToSend;
                consolidationTxIds.push_back(tx.GetHash().ToString());
                numTxCreated++;
            } else {
                LogPrintf("%s: Consolidation transaction FAILED in CommitTransaction, txid=%s\n",opid , tx.GetHash().ToString());

                unlock_notes();
                LogPrint("zrpc", "%s: consolidatoin input notes unlocked\n", getId());

                setConsolidationResult(numTxCreated, amountConsolidated, consolidationTxIds);
                status = false;
                break;
            }

        }
    }

    LogPrintf("%s: Created %d transactions with total Sapling output amount=%s,status=%d\n",opid , numTxCreated, FormatMoney(amountConsolidated), (int)status);
    setConsolidationResult(numTxCreated, amountConsolidated, consolidationTxIds);
    return status;
}

void AsyncRPCOperation_saplingconsolidation::setConsolidationResult(int numTxCreated, const CAmount& amountConsolidated, const std::vector<std::string>& consolidationTxIds) {
    UniValue res(UniValue::VOBJ);
    res.push_back(Pair("num_tx_created", numTxCreated));
    res.push_back(Pair("amount_consolidated", FormatMoney(amountConsolidated)));
    UniValue txIds(UniValue::VARR);
    for (const std::string& txId : consolidationTxIds) {
        txIds.push_back(txId);
    }
    res.push_back(Pair("consolidation_txids", txIds));
    set_result(res);
}

void AsyncRPCOperation_saplingconsolidation::cancel() {
    // Cancelling is how the scheduler stops an in-flight round, so unlike the base
    // class this must be able to move an EXECUTING operation to CANCELLED. What it
    // must not do is overwrite a state that is already terminal: the scheduler
    // cancels the previous operation when it enqueues the next one, and that one may
    // have already SUCCEEDED, whose result would otherwise be relabelled as cancelled.
    if (isSuccess() || isFailed() || isCancelled())
        return;
    set_state(OperationStatus::CANCELLED);
}

UniValue AsyncRPCOperation_saplingconsolidation::getStatus() const {
    UniValue v = AsyncRPCOperation::getStatus();
    UniValue obj = v.get_obj();
    obj.push_back(Pair("method", "saplingconsolidation"));
    obj.push_back(Pair("target_height", targetHeight_));
    return obj;
}

// Lock input notes
void AsyncRPCOperation_saplingconsolidation::lock_notes() {
    LOCK2(cs_main, pwalletMain->cs_wallet);
    for (auto note : z_sapling_inputs_) {
        pwalletMain->LockNote(note.op);
    }
}

// Unlock input notes
void AsyncRPCOperation_saplingconsolidation::unlock_notes() {
    LOCK2(cs_main, pwalletMain->cs_wallet);
    for (auto note : z_sapling_inputs_) {
        pwalletMain->UnlockNote(note.op);
    }
}
