ObsidianDragon/src/wallet/lite_wallet_controller.cpp

// DragonX Wallet - ImGui Edition
// Copyright 2024-2026 The Hush Developers
// Released under the GPLv3

#include "lite_wallet_controller.h"

#include "../data/wallet_state.h"

#include <chrono>
#include <cstdint>
#include <unordered_map>
#include <utility>

#include <nlohmann/json.hpp>
#include <sodium.h>

namespace dragonx {
namespace wallet {

namespace {
constexpr double kZatoshisPerCoin = 100000000.0;  // DRGX has 1e8 zatoshis per coin

// Extract a backend {"error":..} message (string or arbitrary JSON) into a plain string.
std::string extractJsonError(const nlohmann::json& value)
{
    const auto& e = value.at("error");
    return e.is_string() ? e.get<std::string>() : e.dump();
}

// Parse a send/shield response: success is {"txid":".."}, failure is {"error":".."} (NOT an
// "Error:"-prefixed string), and malformed args make the backend return plain-text help.
LiteBroadcastResult parseBroadcastResponse(const LiteBridgeStringResult& bridgeCall)
{
    LiteBroadcastResult out;
    if (!bridgeCall.ok) {
        out.error = bridgeCall.error.empty() ? "backend call failed" : bridgeCall.error;
        return out;
    }
    try {
        const auto j = nlohmann::json::parse(bridgeCall.value);
        if (j.is_object()) {
            if (j.contains("error")) {
                out.error = extractJsonError(j);
                return out;
            }
            if (j.contains("txid") && j.at("txid").is_string()) {
                out.txid = j.at("txid").get<std::string>();
                out.ok = !out.txid.empty();
                if (!out.ok) out.error = "backend returned an empty txid";
                return out;
            }
        }
    } catch (...) {
        // Non-JSON (e.g. the command's plain-text help on bad args) -> generic error below.
    }
    out.error = "could not parse transaction response";
    return out;
}

// Build the JSON-array send payload and broadcast it. litelib_execute passes the whole args
// string as ONE argument (no whitespace splitting), so send MUST use the JSON-array form
// ([{address,amount,memo},..]); the space-separated CLI form would never parse.
LiteBroadcastResult doSend(LiteClientBridge& bridge, const LiteSendRequest& request)
{
    LiteBroadcastResult out;
    if (request.recipients.empty()) {
        out.error = "no recipients";
        return out;
    }
    nlohmann::json arr = nlohmann::json::array();
    for (const auto& r : request.recipients) {
        if (r.address.empty()) {
            out.error = "recipient address is empty";
            return out;
        }
        nlohmann::json o;
        o["address"] = r.address;
        o["amount"] = r.amountZatoshis;  // zatoshis (puposhis)
        if (!r.memo.empty()) o["memo"] = r.memo;
        arr.push_back(std::move(o));
    }
    return parseBroadcastResponse(bridge.execute("send", arr.dump()));
}

LiteBroadcastResult doShield(LiteClientBridge& bridge, const std::string& optionalAddress)
{
    // Empty address -> shield all transparent funds; otherwise shield to the given address.
    return parseBroadcastResponse(bridge.execute("shield", optionalAddress));
}
}  // namespace

void secureWipeLiteSecret(std::string& secret)
{
    if (!secret.empty()) {
        sodium_memzero(&secret[0], secret.size());
    }
    secret.clear();
}

void applyLiteRefreshModelToWalletState(const LiteWalletAppRefreshModel& model,
                                        dragonx::WalletState& state)
{
    if (model.hasBalance) {
        state.privateBalance = static_cast<double>(model.balance.shieldedZatoshis) / kZatoshisPerCoin;
        state.transparentBalance = static_cast<double>(model.balance.transparentZatoshis) / kZatoshisPerCoin;
        state.totalBalance = static_cast<double>(model.balance.totalZatoshis) / kZatoshisPerCoin;
        state.unconfirmedBalance = static_cast<double>(model.balance.unconfirmedZatoshis) / kZatoshisPerCoin;
    }

    if (model.hasAddresses) {
        // Per-address balances from unspent notes/utxos (when the notes command succeeded).
        // Sum the value of confirmed, unspent outputs at each address: skip anything spent,
        // unconfirmed-spent, or still pending (notes/utxos from an unconfirmed received tx) so
        // the per-address figures match the confirmed balance the wallet treats as available.
        std::unordered_map<std::string, std::uint64_t> perAddressZatoshis;
        if (model.hasSpendableOutputs) {
            for (const auto& output : model.spendableOutputs) {
                if (output.spent || output.unconfirmedSpent || output.pending) continue;
                perAddressZatoshis[output.address] += output.valueZatoshis;
            }
        }

        state.addresses.clear();
        state.z_addresses.clear();
        state.t_addresses.clear();
        for (const auto& addr : model.addresses) {
            AddressInfo info;
            info.address = addr.address;
            const auto it = perAddressZatoshis.find(addr.address);
            info.balance = it != perAddressZatoshis.end()
                ? static_cast<double>(it->second) / kZatoshisPerCoin
                : 0.0;
            info.type = (addr.kind == LiteWalletAppAddressKind::Shielded) ? "shielded" : "transparent";
            info.has_spending_key = addr.spendabilityKnown ? addr.spendable : true;
            if (addr.kind == LiteWalletAppAddressKind::Shielded) {
                state.z_addresses.push_back(info);
            } else {
                state.t_addresses.push_back(info);
            }
            state.addresses.push_back(std::move(info));
        }
    }

    if (model.hasTransactions) {
        state.transactions.clear();
        const int64_t chainHeight =
            model.hasSyncStatus ? static_cast<int64_t>(model.sync.chainHeight) : 0;
        for (const auto& record : model.transactions) {
            TransactionInfo tx;
            tx.txid = record.txid;
            if (record.kind == LiteWalletAppTransactionKind::Send) {
                tx.type = "send";
            } else if (record.kind == LiteWalletAppTransactionKind::Receive) {
                tx.type = "receive";
            } else {
                tx.type = record.signedAmountZatoshis < 0 ? "send" : "receive";
            }
            tx.amount = static_cast<double>(record.amountZatoshis) / kZatoshisPerCoin;
            tx.timestamp = record.timestamp;
            tx.address = record.address;
            tx.memo = record.memo;
            if (record.unconfirmed || !record.blockHeight.has_value() || chainHeight == 0) {
                tx.confirmations = record.unconfirmed ? 0 : 1;
            } else {
                const int64_t confs = chainHeight - *record.blockHeight + 1;
                tx.confirmations = confs > 0 ? static_cast<int>(confs) : 0;
            }
            state.transactions.push_back(std::move(tx));
        }
    }

    if (model.hasSyncStatus) {
        state.sync.blocks = static_cast<int>(model.sync.walletHeight);
        state.sync.headers = static_cast<int>(model.sync.chainHeight);
        state.sync.verification_progress = model.sync.progress;
        state.sync.syncing = !model.sync.complete;
    }
}

LiteWalletController::LiteWalletController(WalletCapabilities capabilities,
                                          LiteConnectionSettings connectionSettings,
                                          LiteClientBridge bridge,
                                          LiteWalletControllerOptions options)
    : bridge_(std::make_shared<LiteClientBridge>(std::move(bridge))),
      lifecycle_(capabilities, connectionSettings, bridge_.get(),
                 LiteWalletLifecycleOptions{options.allowBridgeCalls}),
      gateway_(capabilities, connectionSettings, bridge_.get(),
               LiteWalletGatewayOptions{options.allowBridgeCalls}),
      sync_(capabilities, connectionSettings, bridge_.get(),
            LiteSyncServiceOptions{options.allowBridgeCalls})
{
    status_ = lifecycle_.status();
}

LiteWalletController::~LiteWalletController()
{
    stopWorker();  // joins the fast poll worker (short iterations)
    // The sync thread may be blocked in an uninterruptible full scan; detach it. It holds
    // shared refs (bridge_ + syncDone_), so it stays safe and the bridge survives until it
    // finishes — the process is exiting, so a late litelib_shutdown is harmless.
    if (syncThread_.joinable()) syncThread_.detach();
    // Likewise the broadcast thread (send/shield proving): it captures shared refs (bridge +
    // running flag + result slot), never `this`, so detaching is safe.
    if (broadcastThread_.joinable()) broadcastThread_.detach();
}

std::unique_ptr<LiteWalletController> LiteWalletController::createLinked(
    WalletCapabilities capabilities,
    LiteConnectionSettings connectionSettings)
{
    return std::make_unique<LiteWalletController>(
        capabilities,
        std::move(connectionSettings),
        LiteClientBridge::linkedSdxl(),
        LiteWalletControllerOptions{true});
}

void LiteWalletController::onLifecycleResult(const LiteWalletLifecycleResult& result)
{
    status_ = result.status;
    if (result.walletReady) {
        walletOpen_ = true;
        if (persist_) persist_();
        startSync();    // begin background sync on the backend
        startWorker();  // begin periodic refresh -> WalletState (via takeRefreshedModel)
    }
}

void LiteWalletController::startSync()
{
    if (syncLaunched_.exchange(true)) return;
    syncStarted_ = true;
    // The backend `sync` command is a blocking, uninterruptible full chain scan, so run it on
    // a detached thread. Capture shared refs (not the controller) so it is safe to outlive us.
    auto bridge = bridge_;
    auto done = syncDone_;
    syncThread_ = std::thread([bridge, done] {
        if (bridge) bridge->execute("sync", "");  // blocks until synced (or errors out)
        done->store(true);
    });
}

std::optional<LiteWalletAppRefreshModel> LiteWalletController::refreshModel()
{
    if (!walletOpen_.load()) return std::nullopt;

    // syncstatus is fast (reads shared state the sync thread updates). Poll it every time.
    const auto syncResult = sync_.pollSyncStatus(LiteSyncStatusRequest{});

    if (!syncDone_->load()) {
        // Sync still running: publish progress only. Data queries (balance/list) would block
        // until the chain is synced, so don't issue them yet.
        if (!syncResult.ok) return std::nullopt;
        LiteWalletAppRefreshModel model;
        model.hasSyncStatus = true;
        model.sync.walletHeight = syncResult.syncStatus.syncedBlocks;
        model.sync.chainHeight = syncResult.syncStatus.totalBlocks;
        // syncDone_ is authoritative: the detached sync thread is still running, so we are NOT
        // complete regardless of what syncstatus reports. The backend briefly returns the idle
        // shape ({"syncing":"false"} -> parser progress=1.0, complete=true) before the scan
        // starts publishing in-progress status; don't surface that as a misleading 100%/done.
        model.sync.complete = false;
        model.sync.progress = syncResult.syncStatus.complete ? 0.0 : syncResult.syncStatus.progress;
        return model;
    }

    // Synced: full refresh (balance/addresses/transactions are fast now).
    LiteWalletRefreshRequest request;
    if (syncResult.ok) {
        request.haveSyncStatus = true;
        request.syncStatus = syncResult.syncStatus;
    }
    const auto refreshResult = gateway_.refresh(request);
    if (refreshResult.bundle.successfulCommandCount == 0 && !request.haveSyncStatus) {
        return std::nullopt;
    }
    const auto mapped = mapLiteWalletRefreshResult(refreshResult);
    if (!mapped.ok) return std::nullopt;
    return mapped.model;
}

LiteNewAddressResult LiteWalletController::newAddress(bool shielded)
{
    LiteNewAddressResult out;
    if (!walletOpen_.load() || !bridge_) {
        out.error = "no wallet is open";
        return out;
    }
    // Backend address-type tokens: "zs" (shielded) / "R" (transparent) (do_new_address).
    const auto result = bridge_->execute("new", shielded ? "zs" : "R");
    if (!result.ok) {
        out.error = result.error.empty() ? "new address generation failed" : result.error;
        return out;
    }
    // Response is a JSON array with the new address, e.g. ["zs1..."].
    try {
        const auto parsed = nlohmann::json::parse(result.value);
        if (parsed.is_array() && !parsed.empty() && parsed[0].is_string()) {
            out.address = parsed[0].get<std::string>();
        } else if (parsed.is_string()) {
            out.address = parsed.get<std::string>();
        }
    } catch (...) {
        // fall through to the error below
    }
    out.ok = !out.address.empty();
    if (!out.ok) out.error = "could not parse new address response";
    return out;
}

LiteBroadcastResult LiteWalletController::sendTransactionBlocking(const LiteSendRequest& request)
{
    LiteBroadcastResult out;
    if (!walletOpen_.load() || !bridge_) {
        out.error = "no wallet is open";
        return out;
    }
    return doSend(*bridge_, request);
}

LiteBroadcastResult LiteWalletController::shieldFundsBlocking(const std::string& optionalAddress)
{
    LiteBroadcastResult out;
    if (!walletOpen_.load() || !bridge_) {
        out.error = "no wallet is open";
        return out;
    }
    return doShield(*bridge_, optionalAddress);
}

bool LiteWalletController::startBroadcast(std::function<LiteBroadcastResult()> op)
{
    bool expected = false;
    if (!broadcastRunning_->compare_exchange_strong(expected, true)) return false;  // one at a time
    {
        std::lock_guard<std::mutex> lock(*broadcastResultMutex_);
        broadcastResult_->reset();  // drop any already-consumed prior result
    }
    // A previous broadcast thread (already finished) may still be joinable; detach before reuse.
    if (broadcastThread_.joinable()) broadcastThread_.detach();

    auto running = broadcastRunning_;
    auto mutex = broadcastResultMutex_;
    auto slot = broadcastResult_;
    broadcastThread_ = std::thread([op = std::move(op), running, mutex, slot] {
        LiteBroadcastResult result = op();
        {
            std::lock_guard<std::mutex> lock(*mutex);
            *slot = std::move(result);
        }
        running->store(false);
    });
    return true;
}

bool LiteWalletController::sendTransaction(const LiteSendRequest& request)
{
    if (!walletOpen_.load() || !bridge_) return false;
    auto bridge = bridge_;  // shared copy; the op must not capture `this`
    return startBroadcast([bridge, request]() { return doSend(*bridge, request); });
}

bool LiteWalletController::shieldFunds(const std::string& optionalAddress)
{
    if (!walletOpen_.load() || !bridge_) return false;
    auto bridge = bridge_;
    return startBroadcast([bridge, optionalAddress]() { return doShield(*bridge, optionalAddress); });
}

bool LiteWalletController::takeBroadcastResult(LiteBroadcastResult& out)
{
    std::lock_guard<std::mutex> lock(*broadcastResultMutex_);
    if (!broadcastResult_->has_value()) return false;
    out = std::move(**broadcastResult_);
    broadcastResult_->reset();
    return true;
}

LiteImportResult LiteWalletController::importKey(std::string spendingOrViewingKey)
{
    LiteImportResult out;
    if (!walletOpen_.load() || !bridge_) {
        secureWipeLiteSecret(spendingOrViewingKey);
        out.error = "no wallet is open";
        return out;
    }
    if (spendingOrViewingKey.empty()) {
        out.error = "no key provided";
        return out;
    }
    // Transparent WIFs begin with U/5/K/L (TImportCommand); shielded keys begin with
    // "secret-..." / viewing keys "zxview...", so this prefix check won't collide.
    const char first = spendingOrViewingKey[0];
    const bool transparentWif = (first == 'U' || first == '5' || first == 'K' || first == 'L');
    const auto result = bridge_->execute(transparentWif ? "timport" : "import", spendingOrViewingKey);
    secureWipeLiteSecret(spendingOrViewingKey);  // wipe our copy ASAP

    if (!result.ok) {  // do_import_* failures come back "Error:"-prefixed (bridge -> ok=false)
        out.error = result.error.empty() ? "key import failed" : result.error;
        return out;
    }
    try {
        const auto j = nlohmann::json::parse(result.value);
        if (j.is_object() && j.contains("error")) {
            out.error = extractJsonError(j);
            return out;
        }
    } catch (...) {
        // A non-JSON success payload is acceptable; fall through.
    }
    out.detail = result.value;
    out.ok = true;
    return out;
}

LiteExportResult LiteWalletController::exportPrivateKeys(const std::string& optionalAddress)
{
    LiteExportResult out;
    if (!walletOpen_.load() || !bridge_) {
        out.error = "no wallet is open";
        return out;
    }
    // Empty address -> export keys for all addresses; otherwise just the given address.
    const auto result = bridge_->execute("export", optionalAddress);
    if (!result.ok) {
        out.error = result.error.empty() ? "export failed" : result.error;
        return out;
    }
    try {
        const auto j = nlohmann::json::parse(result.value);
        if (j.is_object() && j.contains("error")) {
            out.error = extractJsonError(j);
            return out;
        }
    } catch (...) {
        out.error = "could not parse export response";
        return out;
    }
    out.privateKeysJson = result.value;  // SECRET — caller must not log; wipe after use
    out.ok = true;
    return out;
}

LiteSeedResult LiteWalletController::exportSeed()
{
    LiteSeedResult out;
    if (!walletOpen_.load() || !bridge_) {
        out.error = "no wallet is open";
        return out;
    }
    const auto result = bridge_->execute("seed", "");
    if (!result.ok) {
        out.error = result.error.empty() ? "seed export failed" : result.error;
        return out;
    }
    try {
        const auto j = nlohmann::json::parse(result.value);
        if (j.is_object()) {
            if (j.contains("error")) {
                out.error = extractJsonError(j);
                return out;
            }
            if (j.contains("seed") && j.at("seed").is_string()) {
                out.seedPhrase = j.at("seed").get<std::string>();  // SECRET
                if (j.contains("birthday") && j.at("birthday").is_number_unsigned()) {
                    out.birthday = j.at("birthday").get<std::uint64_t>();
                }
                out.ok = !out.seedPhrase.empty();
                if (!out.ok) out.error = "backend returned an empty seed";
                return out;
            }
        }
    } catch (...) {
        // fall through to the generic parse error
    }
    out.error = "could not parse seed response";
    return out;
}

bool LiteWalletController::refreshWalletState(dragonx::WalletState& state)
{
    auto model = refreshModel();
    if (!model) return false;
    applyLiteRefreshModelToWalletState(*model, state);
    return true;
}

bool LiteWalletController::takeRefreshedModel(LiteWalletAppRefreshModel& out)
{
    std::lock_guard<std::mutex> lock(modelMutex_);
    if (!pendingModel_) return false;
    out = std::move(*pendingModel_);
    pendingModel_.reset();
    return true;
}

void LiteWalletController::startWorker()
{
    if (running_.exchange(true)) return;  // already running
    worker_ = std::thread([this] { workerLoop(); });
}

void LiteWalletController::stopWorker()
{
    if (!running_.exchange(false)) return;  // not running
    wakeCv_.notify_all();
    if (worker_.joinable()) worker_.join();
}

void LiteWalletController::workerLoop()
{
    while (running_.load()) {
        if (walletOpen_.load()) {
            auto model = refreshModel();
            if (model) {
                std::lock_guard<std::mutex> lock(modelMutex_);
                pendingModel_ = std::move(model);
            }
        }
        std::unique_lock<std::mutex> lock(wakeMutex_);
        wakeCv_.wait_for(lock, std::chrono::milliseconds(kRefreshIntervalMs),
                         [this] { return !running_.load(); });
    }
}

LiteWalletLifecycleResult LiteWalletController::createWallet(LiteWalletCreateRequest request)
{
    auto result = lifecycle_.createWallet(request);
    secureWipeLiteSecret(request.passphrase);
    onLifecycleResult(result);
    return result;
}

LiteWalletLifecycleResult LiteWalletController::openWallet(LiteWalletOpenRequest request)
{
    auto result = lifecycle_.openWallet(request);
    secureWipeLiteSecret(request.passphrase);
    onLifecycleResult(result);
    return result;
}

LiteWalletLifecycleResult LiteWalletController::restoreWallet(LiteWalletRestoreRequest request)
{
    auto result = lifecycle_.restoreWallet(request);
    secureWipeLiteSecret(request.seedPhrase);
    secureWipeLiteSecret(request.passphrase);
    onLifecycleResult(result);
    return result;
}

} // namespace wallet
} // namespace dragonx