// Copyright (c) 2016-2024 The Hush developers
// Copyright (c) 2024-2026 The DragonX developers
// Distributed under the GPLv3 software license, see the accompanying
// file COPYING or https://www.gnu.org/licenses/gpl-3.0.en.html
#include "asyncrpcoperation_autoshieldcoinbase.h"
#include "asyncrpcoperation_shieldcoinbase.h" // for ShieldCoinbaseUTXO
#include "consensus/upgrades.h"
#include "hush_defs.h" // ASSETCHAINS_TIMELOCKGTE
#include "init.h"
#include "key_io.h"
#include "main.h"
#include "rpc/protocol.h"
#include "sync.h"
#include "tinyformat.h"
#include "transaction_builder.h"
#include "util.h"
#include "utilmoneystr.h"
#include "wallet.h"

// Sietch dummy zaddr generator (defined in wallet/rpcwallet.cpp)
extern std::string randomSietchZaddr();

// Serialized-size estimates for one spent input (kept in sync with rpcwallet.cpp)
static const size_t AUTOSHIELD_CTXIN_DUST_SIZE = 148;
// Every autoshield tx carries THREE Sapling OutputDescriptions -- the change
// note to destZaddr plus the two Sietch dummies -- at ~948 bytes each. Reserving
// 2000 for "header + sietch outputs" was ~900 bytes short before a single input
// was counted, so a large enough round could build a tx over MAX_TX_SIZE.
static const size_t AUTOSHIELD_SAPLING_OUTPUT_SIZE = 948;
static const size_t AUTOSHIELD_TX_OVERHEAD = (3 * AUTOSHIELD_SAPLING_OUTPUT_SIZE) + 256;
// Hard cap on inputs per round, mirroring z_shieldcoinbase's
// SHIELD_COINBASE_DEFAULT_LIMIT. The byte estimate alone is not a safe bound:
// with a P2PKH coinbase (-mineraddress) the 148-byte figure is exact rather than
// conservative, so an under-estimate translates directly into an oversize tx.
// The remainder is simply shielded on the next round.
static const size_t AUTOSHIELD_MAX_INPUTS = 400;
// Unrelated to the cap above despite sharing the value: this is a SIZE IN BYTES for
// one spent P2SH input, mirroring CTXIN_SPEND_P2SH_SIZE in rpcwallet.cpp.
static const size_t AUTOSHIELD_CTXIN_P2SH_SIZE = 400;

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

AsyncRPCOperation_autoshieldcoinbase::~AsyncRPCOperation_autoshieldcoinbase() {}

void AsyncRPCOperation_autoshieldcoinbase::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->saplingAutoShieldOperationId) {
                pwalletMain->fAutoShieldRunning = 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();

    // ALWAYS advance the interval and clear the running flag, on success AND
    // failure AND exception, so a failed/oversized/locked round still lets the
    // next round fire. Only the CURRENT op does this bookkeeping: if a newer op
    // has already superseded this one, leave its state untouched.
    if (pwalletMain) {
        LOCK2(cs_main, pwalletMain->cs_wallet);
        if (getId() == pwalletMain->saplingAutoShieldOperationId) {
            int tipHeight = (chainActive.Tip() != NULL) ? chainActive.Tip()->GetHeight() : targetHeight_;
            pwalletMain->nextAutoShield    = pwalletMain->autoShieldInterval + tipHeight;
            pwalletMain->fAutoShieldRunning = false;
        }
    }

    set_state(success ? OperationStatus::SUCCESS : OperationStatus::FAILED);
    setResult();

    LogPrintf("%s: autoshield operation finished (status=%s, txs=%d, shielded=%s)\n",
              getId(), getStateAsString(), numTxCreated_, FormatMoney(amountShielded_));
}

// Read-only half of destination resolution, shared with init.cpp so the answer to
// "where will auto-shielding send?" is available before the first round runs rather
// than only after one has fired. Mutates nothing: no key generation, no caching.
AutoShieldDestStatus ResolveAutoShieldDestinationReadOnly(
        libzcash::SaplingPaymentAddress& destOut, std::string& destStrOut, uint32_t& accountOut) {
    accountOut = AUTOSHIELD_ACCOUNT_NONE;

    // 1. Explicit -autoshieldaddress override (validated as a spendable Sapling
    //    zaddr in init.cpp). This doubles as the per-process cache for whatever
    //    the derivation below resolved on an earlier round.
    if (!pwalletMain->autoShieldAddress.empty()) {
        auto decoded = DecodePaymentAddress(pwalletMain->autoShieldAddress);
        if (boost::get<libzcash::SaplingPaymentAddress>(&decoded) != nullptr) {
            destOut    = boost::get<libzcash::SaplingPaymentAddress>(decoded);
            destStrOut = pwalletMain->autoShieldAddress;
            return AutoShieldDestStatus::Resolved;
        }
        return AutoShieldDestStatus::InvalidOverride;
    }

    // 2. Walk the restore window m/32'/coin'/[0, gap)' derived from the seed.
    HDSeed seed;
    if (!pwalletMain->GetHDSeedForDerivation(seed)) {
        return AutoShieldDestStatus::NoSeed;
    }

    // Mirror init.cpp's own clamp, and cap into the hardened index space so
    // (i | ZIP32_HARDENED_KEY_LIMIT) below stays well formed.
    int64_t gapArg = GetArg("-mnemonicsaplinggap", 100);
    if (gapArg < 0) {
        gapArg = 0;
    }
    if (gapArg > (int64_t)ZIP32_HARDENED_KEY_LIMIT) {
        gapArg = (int64_t)ZIP32_HARDENED_KEY_LIMIT;
    }
    const uint32_t saplingGap = (uint32_t)gapArg;

    // Same derivation path as CWallet::GenerateNewSaplingZKey (wallet.cpp:139-152).
    const uint32_t bip44CoinType = Params().BIP44CoinType();
    auto m         = libzcash::SaplingExtendedSpendingKey::Master(seed);
    auto m_32h     = m.Derive(32 | ZIP32_HARDENED_KEY_LIMIT);
    auto m_32h_cth = m_32h.Derive(bip44CoinType | ZIP32_HARDENED_KEY_LIMIT);

    for (uint32_t i = 0; i < saplingGap; i++) {
        auto xsk  = m_32h_cth.Derive(i | ZIP32_HARDENED_KEY_LIMIT);
        auto addr = xsk.DefaultAddress();

        // Spendable AND registered: GetSaplingExtendedSpendingKey resolves
        // addr -> ivk -> fvk -> spending key (keystore.cpp:215-223), so a hit
        // means the wallet both recognises notes sent to `addr` and can spend
        // them. Exactly the pair of properties the shield needs. Lowest index
        // wins: stable for the life of the wallet and reproducible from the seed
        // alone, unlike std::set order over the random diversifier
        // (zcash/Address.hpp:95-98).
        libzcash::SaplingExtendedSpendingKey held;
        if (pwalletMain->GetSaplingExtendedSpendingKey(addr, held)) {
            destOut    = addr;
            destStrOut = EncodePaymentAddress(addr);
            accountOut = i;
            return AutoShieldDestStatus::Resolved;
        }
    }

    return AutoShieldDestStatus::NotFound;
}

// Resolve the Sapling destination for auto-shielded coinbase.
//
// Recoverability is the hard requirement: coinbase we shield must land in an
// address that a bare -mnemonic/-hdseed restore of THIS wallet's seed re-derives
// on its own. A restore pre-derives exactly -mnemonicsaplinggap sapling accounts
// starting at index 0, with saplingAccountCounter reset to 0 (init.cpp:2349-2355),
// so the only self-recoverable destinations are the default addresses of
// m/32'/<coin>'/i' for i < gap.
//
// We therefore DERIVE those accounts from the seed and pick the lowest index the
// wallet already holds. Deriving is the only authoritative test. In particular
// CKeyMetadata is NOT evidence of provenance: z_importkey / z_importwallet copy
// both hdKeypath and seedFp verbatim out of the import source
// (wallet.cpp:5522-5529 <- rpcdump.cpp:511-516), so a foreign key can claim any
// keypath and any seed fingerprint. Filtering on metadata would let an imported
// key win as "account 0" and silently receive every shielded reward.
//
// Caller must hold cs_wallet and must already have checked the wallet is unlocked.
bool AsyncRPCOperation_autoshieldcoinbase::resolveDestination(
        libzcash::SaplingPaymentAddress& destOut, std::string& destStrOut) {

    uint32_t account = AUTOSHIELD_ACCOUNT_NONE;
    switch (ResolveAutoShieldDestinationReadOnly(destOut, destStrOut, account)) {
    case AutoShieldDestStatus::Resolved:
        // Cache for the life of the process; the override branch of the resolver
        // short-circuits later rounds. Safe: we only cache post-validation.
        pwalletMain->autoShieldAddress = destStrOut;
        if (account == AUTOSHIELD_ACCOUNT_NONE) {
            LogPrintf("%s: autoshield destination %s (configured)\n", getId(), destStrOut);
        } else {
            LogPrintf("%s: autoshield destination %s (seed-derived sapling account %u)\n",
                      getId(), destStrOut, (unsigned)account);
        }
        return true;
    case AutoShieldDestStatus::InvalidOverride:
        LogPrintf("%s: configured -autoshieldaddress is not a valid Sapling address\n", getId());
        return false;
    case AutoShieldDestStatus::NoSeed:
        LogPrintf("%s: no HD seed available; refusing to pick an autoshield destination\n", getId());
        return false;
    case AutoShieldDestStatus::NotFound:
        break; // nothing in the window yet: fall through and derive one
    }

    // Re-establish the derivation context the resolver used, for step 3 below.
    HDSeed seed;
    if (!pwalletMain->GetHDSeedForDerivation(seed)) {
        LogPrintf("%s: no HD seed available; refusing to pick an autoshield destination\n", getId());
        return false;
    }
    int64_t gapArg = GetArg("-mnemonicsaplinggap", 100);
    if (gapArg < 0) {
        gapArg = 0;
    }
    if (gapArg > (int64_t)ZIP32_HARDENED_KEY_LIMIT) {
        gapArg = (int64_t)ZIP32_HARDENED_KEY_LIMIT;
    }
    const uint32_t saplingGap    = (uint32_t)gapArg;
    const uint32_t bip44CoinType = Params().BIP44CoinType();
    auto m         = libzcash::SaplingExtendedSpendingKey::Master(seed);
    auto m_32h     = m.Derive(32 | ZIP32_HARDENED_KEY_LIMIT);
    auto m_32h_cth = m_32h.Derive(bip44CoinType | ZIP32_HARDENED_KEY_LIMIT);

    // 3. Nothing usable in the window yet: derive the next account, but only if
    //    GenerateNewSaplingZKey will land INSIDE the window. It does NOT derive
    //    at saplingAccountCounter: its do/while skips every index whose spending
    //    key we already hold (wallet.cpp:150-157), so a bare "counter < gap"
    //    test is unsound - counter 98 with gap 100 can still land on 100.
    //    Predict min{ i >= counter : we do not hold i } instead.
    const uint32_t counter = pwalletMain->GetHDChain().saplingAccountCounter;
    uint32_t predicted     = saplingGap; // sentinel: "would land outside the window"
    for (uint32_t i = counter; i < saplingGap; i++) {
        auto xsk = m_32h_cth.Derive(i | ZIP32_HARDENED_KEY_LIMIT);
        if (!pwalletMain->HaveSaplingSpendingKey(xsk.expsk.full_viewing_key())) {
            predicted = i;
            break;
        }
    }

    if (predicted == saplingGap) {
        LogPrintf("%s: no free sapling account below -mnemonicsaplinggap=%u (account counter is %u). "
                  "A newly derived z-address would NOT be re-derived by a seed restore, so the "
                  "shielded coinbase could not be recovered from the seed alone. Refusing to "
                  "autoshield this round. Fix: point -autoshieldaddress at an existing in-gap "
                  "wallet z-address, or raise -mnemonicsaplinggap here AND use the same value on "
                  "any future restore.\n",
                  getId(), (unsigned)saplingGap, (unsigned)counter);
        return false;
    }

    if (pwalletMain->IsLocked()) {
        LogPrintf("%s: wallet is locked; cannot derive an autoshield destination z-address\n", getId());
        return false;
    }

    try {
        auto expectedAddr = m_32h_cth.Derive(predicted | ZIP32_HARDENED_KEY_LIMIT).DefaultAddress();
        libzcash::SaplingPaymentAddress newAddr = pwalletMain->GenerateNewSaplingZKey();

        // Post-verify rather than trust the prediction: cheap, and it closes the
        // whole class of "the counter moved further than expected" bugs.
        if (!(newAddr == expectedAddr)) {
            LogPrintf("%s: newly derived z-address is not sapling account %u as predicted "
                      "(counter %u -> %u); not using it as the autoshield destination\n",
                      getId(), (unsigned)predicted, (unsigned)counter,
                      (unsigned)pwalletMain->GetHDChain().saplingAccountCounter);
            return false;
        }

        destOut    = newAddr;
        destStrOut = EncodePaymentAddress(newAddr);
        pwalletMain->autoShieldAddress = destStrOut;
        LogPrintf("%s: generated new autoshield destination z-address %s (seed-derived sapling account %u)\n",
                  getId(), destStrOut, (unsigned)predicted);
        return true;
    } catch (const std::exception& e) {
        LogPrintf("%s: could not generate a destination z-address: %s\n", getId(), e.what());
        return false;
    }
}

bool AsyncRPCOperation_autoshieldcoinbase::main_impl() {
    auto opid = getId();
    LogPrintf("%s: Beginning asyncrpcoperation_autoshieldcoinbase.\n", opid);
    auto consensusParams = Params().GetConsensus();

    int tipHeight;
    {
        LOCK(cs_main);
        tipHeight = (chainActive.Tip() != NULL) ? chainActive.Tip()->GetHeight() : targetHeight_;
    }

    // Don't create a tx that could be mined before, but expire after, a NU
    // activation. Key this off tipHeight (the height we actually set the expiry
    // from below), not the stale enqueue-time targetHeight_, so a queue delay
    // cannot slip a straddling expiry past this guard.
    auto nextActivationHeight = NextActivationHeight(tipHeight, consensusParams);
    if (nextActivationHeight && tipHeight + AUTO_OP_EXPIRY_DELTA >= nextActivationHeight.get()) {
        LogPrintf("%s: autoshield tx could expire across a NU activation. Skipping this round.\n", opid);
        return true;
    }

    libzcash::SaplingPaymentAddress destZaddr;
    std::string destStr;
    std::vector<ShieldCoinbaseUTXO> inputs;

    // Proof building below runs WITHOUT cs_wallet (deliberately, so wallet RPCs
    // are not stalled), which leaves a multi-second window in which a manual
    // z_shieldcoinbase or z_sendmany over the same miner address would re-select
    // these same coinbase outputs. AvailableCoins honours IsLockedCoin, so lock
    // them for the duration exactly as z_shieldcoinbase does. RAII because there
    // are several early returns between here and commit, and a leaked lock would
    // silently exclude those coins from every future round.
    struct ScopedCoinLocks {
        std::vector<COutPoint> locked;
        ~ScopedCoinLocks() {
            // A destructor is noexcept by default; letting the lock acquisition
            // escape would turn a contended mutex into std::terminate.
            try {
                if (locked.empty()) return;
                LOCK2(cs_main, pwalletMain->cs_wallet);
                // UnlockCoin takes a non-const reference (upstream signature).
                for (COutPoint& op : locked) pwalletMain->UnlockCoin(op);
            } catch (...) {}
        }
    } coinLocks;
    CAmount shieldedValue      = 0;
    unsigned int max_tx_size   = MAX_TX_SIZE_AFTER_SAPLING;

    {
        LOCK2(cs_main, pwalletMain->cs_wallet);

        // Defensive: the scheduler already skips while locked, but the wallet
        // could have been locked between enqueue and execution.
        if (pwalletMain->IsLocked()) {
            LogPrintf("%s: wallet is locked, skipping autoshield round\n", opid);
            return true;
        }

        if (!resolveDestination(destZaddr, destStr)) {
            LogPrintf("%s: no spendable destination z-address available, skipping\n", opid);
            return true;
        }

        // Gather matured, spendable coinbase UTXOs, byte-capped to a single tx.
        // AvailableCoins excludes immature coinbase unconditionally (wallet.cpp,
        // `IsCoinBase() && GetBlocksToMaturity() > 0`) and only ever returns outputs
        // we own, so external -mineraddress / pool coinbase yields zero inputs. The
        // second argument here is fOnlyConfirmed, not fOnlySpendable.
        size_t estimatedTxSize = AUTOSHIELD_TX_OVERHEAD;
        std::vector<COutput> vecOutputs;
        pwalletMain->AvailableCoins(vecOutputs, true, NULL, false, true);
        for (const COutput& out : vecOutputs) {
            if (!out.fSpendable || !out.tx->IsCoinBase()) {
                continue;
            }
            CTxDestination address;
            if (!ExtractDestination(out.tx->vout[out.i].scriptPubKey, address)) {
                continue;
            }
            size_t increase = (boost::get<CScriptID>(&address) != nullptr)
                                  ? AUTOSHIELD_CTXIN_P2SH_SIZE : AUTOSHIELD_CTXIN_DUST_SIZE;
            if (inputs.size() >= AUTOSHIELD_MAX_INPUTS) {
                LogPrintf("%s: reached per-round input cap (%d); deferring remaining coinbase to next round\n",
                          opid, (int)AUTOSHIELD_MAX_INPUTS);
                break;
            }
            if (estimatedTxSize + increase >= max_tx_size) {
                // Size-safe batch; the remainder is shielded next round.
                LogPrintf("%s: reached per-tx size cap; deferring remaining coinbase to next round\n", opid);
                break;
            }
            estimatedTxSize += increase;
            ShieldCoinbaseUTXO utxo = { out.tx->GetHash(), out.i,
                                        out.tx->vout[out.i].scriptPubKey,
                                        out.tx->vout[out.i].nValue };
            inputs.push_back(utxo);
            shieldedValue += out.tx->vout[out.i].nValue;
        }

        for (const ShieldCoinbaseUTXO& t : inputs) {
            COutPoint outpt(t.txid, t.vout);
            pwalletMain->LockCoin(outpt);
            coinLocks.locked.push_back(outpt);
        }
    }

    CAmount fee = pwalletMain->autoShieldFee;

    if (inputs.size() < (size_t)pwalletMain->autoShieldMinUtxos) {
        LogPrintf("%s: %d matured coinbase utxo(s) < min %d, skipping this round\n",
                  opid, (int)inputs.size(), pwalletMain->autoShieldMinUtxos);
        return true;
    }
    if (shieldedValue <= fee) {
        LogPrintf("%s: matured coinbase value %s <= fee %s, skipping\n",
                  opid, FormatMoney(shieldedValue), FormatMoney(fee));
        return true;
    }

    // Common outgoing viewing key derived from the HD seed, exactly as
    // z_shieldcoinbase does for t->z (keeps the note recoverable).
    HDSeed seed;
    if (!pwalletMain->GetHDSeedForDerivation(seed)) {
        LogPrintf("%s: HD seed not available, skipping\n", opid);
        return true;
    }
    uint256 ovk = ovkForShieldingFromTaddr(seed);

    // Build the t->z shield tx. Proof generation happens in Build() WITHOUT
    // holding cs_wallet (mirrors the sweep op) so we don't stall wallet RPCs.
    // tipHeight, not targetHeight_: the builder's height selects the consensus
    // branch id (transaction_builder.cpp CurrentEpochBranchId), and the NU-straddle
    // guard above plus SetExpiryHeight below are both keyed off tipHeight. Using the
    // stale enqueue-time height here meant the guard was checking a height the
    // transaction was not actually signed against.
    auto builder = TransactionBuilder(consensusParams, tipHeight, pwalletMain);
    builder.SetExpiryHeight(tipHeight + AUTO_OP_EXPIRY_DELTA);
    builder.SetFee(fee);

    for (const auto& t : inputs) {
        if (t.amount >= ASSETCHAINS_TIMELOCKGTE) {
            builder.SetLockTime((uint32_t)tipHeight);
            builder.AddTransparentInput(COutPoint(t.txid, t.vout), t.scriptPubKey, t.amount, 0xfffffffe);
        } else {
            builder.AddTransparentInput(COutPoint(t.txid, t.vout), t.scriptPubKey, t.amount);
        }
    }

    // All input value (less fee) goes back to our own z-address as change.
    builder.SendChangeTo(destZaddr, ovk);

    // Sietch padding: mirror z_shieldcoinbase's two dummy zouts so autoshield
    // txs are structurally indistinguishable from manual coinbase shields.
    for (int i = 0; i < 2; i++) {
        auto zdust = DecodePaymentAddress(randomSietchZaddr());
        if (IsValidPaymentAddress(zdust)) {
            builder.AddSaplingOutput(ovk, boost::get<libzcash::SaplingPaymentAddress>(zdust), 0);
        }
    }

    auto maybe_tx = builder.Build();
    if (!maybe_tx) {
        LogPrintf("%s: Failed to build autoshield transaction.\n", opid);
        return false;
    }
    CTransaction tx = maybe_tx.get();

    if (isCancelled()) {
        LogPrintf("%s: Cancelled before commit.\n", opid);
        return false;
    }

    if (pwalletMain->CommitAutomatedTx(tx)) {
        LogPrintf("%s: shielded %s coinbase (%d utxos) into %s via txid=%s\n",
                  opid, FormatMoney(shieldedValue - fee), (int)inputs.size(),
                  destStr, tx.GetHash().ToString());
        amountShielded_ += shieldedValue - fee;
        shieldTxIds_.push_back(tx.GetHash().ToString());
        numTxCreated_++;
        return true;
    }

    LogPrintf("%s: autoshield tx FAILED in CommitTransaction, txid=%s\n", opid, tx.GetHash().ToString());
    return false;
}

void AsyncRPCOperation_autoshieldcoinbase::setResult() {
    UniValue res(UniValue::VOBJ);
    res.push_back(Pair("num_tx_created", numTxCreated_));
    res.push_back(Pair("amount_shielded", FormatMoney(amountShielded_)));
    UniValue txIds(UniValue::VARR);
    for (const std::string& txId : shieldTxIds_) {
        txIds.push_back(txId);
    }
    res.push_back(Pair("shield_txids", txIds));
    set_result(res);
}

void AsyncRPCOperation_autoshieldcoinbase::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_autoshieldcoinbase::getStatus() const {
    UniValue v = AsyncRPCOperation::getStatus();
    UniValue obj = v.get_obj();
    obj.push_back(Pair("method", "autoshieldcoinbase"));
    obj.push_back(Pair("target_height", targetHeight_));
    return obj;
}
