// 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;
static const size_t AUTOSHIELD_CTXIN_P2SH_SIZE = 400;
// Expire unmined autoshield txs after this many blocks, so a tx cannot straddle
// a network-upgrade activation.
static const int AUTOSHIELD_EXPIRY_DELTA = 15;

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 (const UniValue& objError) {
        int code = find_value(objError, "code").get_int();
        std::string message = find_value(objError, "message").get_str();
        set_error_code(code);
        set_error_message(message);
    } catch (const runtime_error& e) {
        set_error_code(-1);
        set_error_message("runtime error: " + string(e.what()));
    } catch (const logic_error& e) {
        set_error_code(-1);
        set_error_message("logic error: " + string(e.what()));
    } catch (const exception& e) {
        set_error_code(-1);
        set_error_message("general exception: " + string(e.what()));
    } catch (...) {
        set_error_code(-2);
        set_error_message("unknown error");
    }

    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_));
}

// 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) {

    // 1. Explicit -autoshieldaddress override (validated as a spendable Sapling
    //    zaddr at init.cpp:2494-2506). This also serves as the per-process cache
    //    for whatever step 2/3 resolved.
    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 true;
        }
        LogPrintf("%s: configured -autoshieldaddress is not a valid Sapling address\n", getId());
        return false;
    }

    // 2. Walk the restore window m/32'/coin'/[0, gap)' derived from the seed.
    HDSeed seed;
    if (!pwalletMain->GetHDSeedForDerivation(seed)) {
        LogPrintf("%s: no HD seed available; refusing to pick an autoshield destination\n", getId());
        return false;
    }

    // Mirror init.cpp:2349-2350'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);
            // Cache for the life of the process; step 1 short-circuits later
            // rounds. Safe: we only cache post-validation.
            pwalletMain->autoShieldAddress = destStrOut;
            LogPrintf("%s: autoshield destination %s (seed-derived sapling account %u, gap %u)\n",
                      getId(), destStrOut, (unsigned)i, (unsigned)saplingGap);
            return true;
        }
    }

    // 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 + AUTOSHIELD_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;
    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 with fOnlySpendable already excludes immature coinbase
        // (< COINBASE_MATURITY) and outputs we don't own, so external
        // -mineraddress / pool coinbase naturally yields zero inputs.
        size_t estimatedTxSize = 2000; // header + sietch outputs headroom
        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 (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;
        }
    }

    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.
    auto builder = TransactionBuilder(consensusParams, targetHeight_, pwalletMain);
    builder.SetExpiryHeight(tipHeight + AUTOSHIELD_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() {
    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;
}