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

#include "wallet_scan_governor.h"

#include <algorithm>
#include <limits>

namespace dragonx {
namespace services {

ScanVisibility scanVisibilityForPage(ui::NavPage page, bool windowBackground)
{
    ScanVisibility v;
    if (windowBackground) return v;
    using NP = ui::NavPage;
    switch (page) {
    case NP::Overview:
        v.balance = v.addresses = v.history = true;   // totals, per-address balances, recent activity
        break;
    case NP::Send:
    case NP::Receive:
        v.addresses = v.history = true;               // per-address balances, recent sends / receives
        break;
    case NP::History:
        v.history = true;
        break;
    case NP::Market:
        v.balance = true;                             // portfolio value
        break;
    case NP::Chat:
        v.chat = true;
        break;
    default:
        break;
    }
    return v;
}

double hiddenBalanceIntervalSec(double emaCostSec, bool autoShieldArmed)
{
    const double floorSec = autoShieldArmed ? 120.0 : 600.0;
    return std::max(floorSec, 50.0 * std::max(0.0, emaCostSec));
}

// Minimum gap between VISIBLE balance scans, as a multiple of the measured scan cost. A whole-wallet
// z_gettotalbalance / z_listunspent holds the daemon's cs_main for its whole run (tens of seconds on a
// very large note set); re-polling before it recovers starves block connection and times the next scan
// out, which leaves per-address balances at 0. 3x the cost gives the node room to breathe while staying
// responsive — on a normal wallet the cost is milliseconds, so this floor is effectively zero.
constexpr double kVisibleBalanceCostMultiple = 3.0;

BalancePollDecision decideBalancePoll(const BalancePollInputs& in)
{
    if (!in.allowBalance || !in.gateOpen) return BalancePollDecision::Skip;
    // The session's first read runs even while the node catches up, once per connection and only when
    // someone can see it: otherwise the total reads 0 for the whole catch-up. After that (or if that one
    // attempt failed) the balance waits for the sync to finish, a forced one too.
    if (in.neverFetched && !(in.syncing && (in.background || in.firstReadTried))) return BalancePollDecision::Explicit;
    if (in.syncing) return BalancePollDecision::Skip;
    // The user's own send, now committed. In the background it waits for the window (the routine rules below
    // still apply, e.g. the auto-shield poll); coming back forces it.
    if (in.forced && !in.background) return BalancePollDecision::Explicit;
    if (!in.changed) return BalancePollDecision::Skip;
    if (in.visible) {
        // Cost-proportional floor: don't re-scan before ~3x the measured scan cost has elapsed, so a
        // wallet whose whole-wallet sweep costs many seconds can't saturate the node by re-polling before
        // it recovers. Negligible for small wallets. The forced / first-read paths above bypass this.
        if (in.sinceLastScanSec < kVisibleBalanceCostMultiple * std::max(0.0, in.emaCostSec))
            return BalancePollDecision::Skip;
        // Just came into view with an old value: fetch it now rather than queue behind the other streams.
        if (!in.waiting && in.shownAgeSec >= in.visibleFloorSec && in.sinceLastScanSec >= in.visibleFloorSec)
            return BalancePollDecision::Explicit;
        return BalancePollDecision::AskGovernor;
    }
    if (in.sinceLastScanSec < hiddenBalanceIntervalSec(in.emaCostSec, in.autoShieldArmed))
        return BalancePollDecision::Skip;
    return BalancePollDecision::AskGovernor;
}

double WalletScanGovernor::budgetForTipLag(double base, int blocksBehind)
{
    if (blocksBehind <= 0) return base;
    if (blocksBehind == 1) return base * 2.0 / 3.0;
    return base / 3.0;
}

bool WalletScanGovernor::signalsMatch(ScanStream stream, const ScanSignals& current) const
{
    const Stream& st = at(stream);
    if (!st.basisValid) return false;
    const unsigned mask = cfg_.streams[static_cast<std::size_t>(stream)].signalMask;
    if (mask & kScanSignalTip) {
        if (current.tipHash.empty() || st.basis.tipHash.empty() || current.tipHash != st.basis.tipHash) return false;
    }
    if (mask & kScanSignalMempool) {
        if (!current.mempoolKnown || !st.basis.mempoolKnown) return false;
        if (current.mempoolTxs != st.basis.mempoolTxs || current.mempoolBytes != st.basis.mempoolBytes) return false;
    }
    if (mask & kScanSignalEpoch) {
        if (current.walletEpoch != st.basis.walletEpoch) return false;
    }
    return true;
}

bool WalletScanGovernor::needed(ScanStream stream, const ScanSignals& current, double now) const
{
    const Stream& st = at(stream);
    if (!st.basisValid) return true;
    if (!signalsMatch(stream, current)) return true;
    const double maxAge = cfg_.streams[static_cast<std::size_t>(stream)].maxAgeSec;
    return maxAge > 0.0 && now - st.basisAt >= maxAge;
}

double WalletScanGovernor::floorSec(const Stream& stream) const
{
    return std::max(stream.askVisible ? cfg_.visibleFloorSec : cfg_.hiddenFloorSec,
                    cfg_.floorCostMultiple * stream.emaCostMs / 1000.0);
}

ScanStream WalletScanGovernor::bestWaiter(double now) const
{
    const bool quiet = now < quietUntil_;
    ScanStream best = ScanStream::Count;
    double bestKey = std::numeric_limits<double>::infinity();
    for (std::size_t i = 0; i < streams_.size(); ++i) {
        const Stream& st = streams_[i];
        if (st.waitingSince < 0.0) continue;
        if (now - st.lastAskAt > cfg_.waiterTtlSec) continue;            // stopped asking
        if (quiet && now - st.waitingSince < floorSec(st)) continue;     // during quiet, only via its floor
        const double key = st.waitingSince - (st.askVisible ? cfg_.visibleHeadStartSec : 0.0);
        if (key < bestKey) {
            bestKey = key;
            best = static_cast<ScanStream>(i);
        }
    }
    return best;
}

ScanTicket WalletScanGovernor::issue(ScanStream stream, ScanKind kind, const ScanSignals& current, double now,
                                     bool servesRoutine)
{
    ScanTicket t;
    t.stream = stream;
    t.kind = kind;
    t.id = ++nextId_;
    t.generation = generation_;
    t.dispatchedAt = now;
    t.servesRoutine = servesRoutine;
    t.basis = current;
    live_.push_back(Live{t.id, stream, kind, now, servesRoutine});
    return t;
}

ScanTicket WalletScanGovernor::ticketFor(const Live& live) const
{
    ScanTicket t;
    t.stream = live.stream;
    t.kind = live.kind;
    t.id = live.id;
    t.generation = generation_;
    t.dispatchedAt = live.dispatchedAt;
    t.servesRoutine = live.servesRoutine;
    return t;
}

ScanTicket WalletScanGovernor::tryAdmit(ScanStream stream, bool visible, const ScanSignals& current, double now)
{
    Stream& st = at(stream);
    if (st.waitingSince < 0.0 || now - st.lastAskAt > cfg_.waitLapseSec) st.waitingSince = now;   // a new wait
    st.lastAskAt = now;
    st.askVisible = visible;
    pruneLeases(now);
    if (!live_.empty()) return {};                    // one wallet scan in the daemon at a time
    if (bestWaiter(now) != stream) return {};
    ScanTicket t = issue(stream, ScanKind::Routine, current, now, true);
    t.waitedSince = st.waitingSince;
    st.waitingSince = -1.0;
    return t;
}

ScanTicket WalletScanGovernor::beginExplicit(ScanStream stream, const ScanSignals& current, double now,
                                             bool servesRoutine)
{
    ScanTicket t = issue(stream, ScanKind::Explicit, current, now, servesRoutine);
    if (servesRoutine) {
        Stream& st = at(stream);
        t.waitedSince = st.waitingSince;
        st.waitingSince = -1.0;   // served
    }
    return t;
}

void WalletScanGovernor::eraseLive(std::uint64_t id)
{
    live_.erase(std::remove_if(live_.begin(), live_.end(), [id](const Live& l) { return l.id == id; }), live_.end());
}

namespace {
// A ticket that sent nothing gives back the wait it ended (keeping a wait begun since, if older).
void restoreWait(double& waitingSince, double waitedSince)
{
    if (waitedSince < 0.0) return;
    waitingSince = waitingSince < 0.0 ? waitedSince : std::min(waitingSince, waitedSince);
}
}  // namespace

void WalletScanGovernor::cancel(const ScanTicket& ticket)
{
    if (!ticket || ticket.generation != generation_) return;
    const bool wasLive = std::any_of(live_.begin(), live_.end(), [&](const Live& l) { return l.id == ticket.id; });
    if (!wasLive) return;
    eraseLive(ticket.id);
    restoreWait(at(ticket.stream).waitingSince, ticket.waitedSince);
}

void WalletScanGovernor::complete(const ScanTicket& ticket, double costMs, ScanOutcome outcome, double now)
{
    if (!ticket || ticket.generation != generation_) return;   // from before a reset()
    const bool wasLive = std::any_of(live_.begin(), live_.end(), [&](const Live& l) { return l.id == ticket.id; });
    if (!wasLive) return;                                       // already completed (or dropped as an orphan)
    eraseLive(ticket.id);
    Stream& st = at(ticket.stream);
    if (outcome == ScanOutcome::Skipped) {   // nothing was sent: no charge, and it keeps its place in line
        restoreWait(st.waitingSince, ticket.waitedSince);
        return;
    }
    st.lastScanAt = std::max(st.lastScanAt, ticket.dispatchedAt);
    double ms = std::max(0.0, costMs);
    if (outcome == ScanOutcome::TimedOut && ticket.servesRoutine) ms = std::max(ms, st.emaCostMs);
    if (ms > 0.0) {
        // Quiet long enough that this scan was `budget` of the time since it started.
        const double debt = ms / 1000.0 * (1.0 / budgetForTipLag(cfg_.budget, tipLag_) - 1.0);
        const double cap = now + (ticket.kind == ScanKind::Explicit ? cfg_.explicitDebtCapSec : cfg_.maxQuietSec);
        quietUntil_ = std::max(quietUntil_, std::min(std::max(quietUntil_, now) + debt, cap));
    }
    if (!ticket.servesRoutine) return;   // charged; the routine scans' cost and basis are theirs
    if (outcome == ScanOutcome::Ok || outcome == ScanOutcome::Partial || outcome == ScanOutcome::TimedOut) {
        // Rises at once, decays slowly: a scan that got expensive is spaced out right away.
        st.emaCostMs = (st.emaCostMs <= 0.0 || ms > st.emaCostMs) ? ms : 0.7 * st.emaCostMs + 0.3 * ms;
    }
    if (outcome == ScanOutcome::Failed || outcome == ScanOutcome::TimedOut) {
        if (st.pendingValid) st.pendingPoisoned = true;
    } else if (outcome == ScanOutcome::Partial) {
        if (!st.pendingValid) {
            st.pending = ticket.basis;
            st.pendingAt = ticket.dispatchedAt;
            st.pendingValid = true;
            st.pendingPoisoned = false;
        }
    } else if (outcome == ScanOutcome::Ok) {
        // A pass that started at the pending basis is current for that basis (the earliest), unless one of
        // its dispatches failed.
        if (!(st.pendingValid && st.pendingPoisoned)) {
            st.basis = st.pendingValid ? st.pending : ticket.basis;
            st.basisAt = st.pendingValid ? st.pendingAt : ticket.dispatchedAt;
            st.basisValid = true;
        }
        st.pendingValid = false;
        st.pendingPoisoned = false;
    }
}

void WalletScanGovernor::dropOrphans(ScanStream stream, double now)
{
    std::vector<Live> orphans;
    for (const auto& l : live_) {
        if (l.stream == stream) orphans.push_back(l);
    }
    for (const auto& l : orphans) complete(ticketFor(l), 0.0, ScanOutcome::TimedOut, now);
}

void WalletScanGovernor::pruneLeases(double now)
{
    std::vector<Live> expired;
    for (const auto& l : live_) {
        const double lease = std::max(cfg_.minLeaseSec, 4.0 * at(l.stream).emaCostMs / 1000.0);
        if (now - l.dispatchedAt > lease) expired.push_back(l);
    }
    for (const auto& l : expired) complete(ticketFor(l), 0.0, ScanOutcome::TimedOut, now);
}

void WalletScanGovernor::invalidate(ScanStream stream)
{
    Stream& st = at(stream);
    st.basisValid = false;
    st.pendingValid = false;
    st.pendingPoisoned = false;
}

void WalletScanGovernor::reset()
{
    const std::uint64_t nextGeneration = generation_ + 1;
    const std::uint64_t nextId = nextId_;
    *this = WalletScanGovernor(cfg_);
    generation_ = nextGeneration;
    nextId_ = nextId;
}

bool WalletScanGovernor::inFlight(ScanStream stream) const
{
    return std::any_of(live_.begin(), live_.end(), [stream](const Live& l) { return l.stream == stream; });
}

double WalletScanGovernor::secondsSinceScan(ScanStream stream, double now) const
{
    const double at_ = at(stream).lastScanAt;
    return at_ < 0.0 ? 1e9 : now - at_;
}

bool WalletScanGovernor::waiting(ScanStream stream, double now) const
{
    const Stream& st = at(stream);
    return st.waitingSince >= 0.0 && now - st.lastAskAt <= cfg_.waitLapseSec;
}

bool WalletScanGovernor::heldWithinFloor(ScanStream stream, double now) const
{
    if (!waiting(stream, now)) return false;
    const Stream& st = at(stream);
    return now - st.waitingSince < floorSec(st) + st.emaCostMs / 1000.0;
}

double WalletScanGovernor::visibleFloorSecFor(ScanStream stream) const
{
    return std::max(cfg_.visibleFloorSec, cfg_.floorCostMultiple * at(stream).emaCostMs / 1000.0);
}

}  // namespace services
}  // namespace dragonx