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

#include "shield_dialog.h"
#include "../../app.h"
#include "../../config/version.h"
#include "../../data/wallet_state.h"
#include "../../rpc/rpc_client.h"
#include "../../rpc/rpc_worker.h"
#include "../../util/i18n.h"
#include "../notifications.h"
#include "../schema/ui_schema.h"
#include "../material/draw_helpers.h"
#include "imgui.h"

#include <vector>
#include <string>
#include <cstdio>

namespace dragonx {
namespace ui {

// ── Static dialog state ─────────────────────────────────────────────────────────────────────────
static bool   s_open = false;
static ShieldDialog::Mode s_mode = ShieldDialog::Mode::ShieldCoinbase;
static bool   s_consolidate = false;              // opened from the wallet-bloat nudge (shielded preset + framing)
static int    s_src = 2;                          // merge source: 0 = transparent, 1 = shielded, 2 = both
static char   s_from_address[512] = "*";
static char   s_to_address[512] = "";
static int    s_selected_zaddr_idx = -1;
static double s_fee = DRAGONX_DEFAULT_FEE;
static int    s_utxo_limit = 50;                  // overridden by schema at runtime
static bool   s_advanced = false;                 // Advanced (fee + batch size) disclosure
static bool   s_confirm = false;                  // inline "confirm before moving funds" phase
static bool   s_operation_pending = false;
static bool   s_op_terminal = false;              // async op reached success/failed — freeze inputs
static std::string s_operation_id;
static std::string s_status_message;
static double s_last_poll = 0.0;                  // live-progress self-poll timer (ImGui::GetTime seconds)
// Scope of what can be consolidated (fetched once on open, merge mode only).
static bool   s_scope_loading = false;
static bool   s_scope_loaded = false;
static int    s_t_count = 0, s_z_count = 0;
static double s_t_amount = 0.0, s_z_amount = 0.0;
static bool   s_creating_addr = false;            // z_getnewaddress in flight (empty state)

static void resetTransient()
{
    s_operation_pending = false;
    s_op_terminal = false;
    s_confirm = false;
    s_status_message.clear();
    s_operation_id.clear();
    s_creating_addr = false;
}

// Count + sum spendable transparent UTXOs and shielded notes so the user can see the scope of a
// consolidation (and how many batches it may take). Read-only; runs off the UI thread.
static void loadScope(App* app)
{
    if (!app || !app->worker()) return;
    s_scope_loading = true; s_scope_loaded = false;
    app->worker()->post([app, rpc = app->rpc()]() -> rpc::RPCWorker::MainCb {
        int tC = 0, zC = 0; double tA = 0.0, zA = 0.0; std::string error;
        try {
            rpc::RPCClient::TraceScope trace("Shield dialog / Scope count");
            nlohmann::json us = rpc->call("listunspent", nlohmann::json::array({0}));
            if (us.is_array()) for (const auto& u : us) {
                if (u.value("confirmations", 0) >= 1 && u.value("spendable", true)) { ++tC; tA += u.value("amount", 0.0); }
            }
            nlohmann::json zs = rpc->call("z_listunspent", nlohmann::json::array({0}));
            if (zs.is_array()) for (const auto& z : zs) {
                if (z.value("confirmations", 0) >= 1) { ++zC; zA += z.value("amount", 0.0); }
            }
        } catch (const std::exception& e) { error = e.what(); }
        return [tC, zC, tA, zA, error]() {
            s_scope_loading = false; s_scope_loaded = error.empty();
            s_t_count = tC; s_z_count = zC; s_t_amount = tA; s_z_amount = zA;
            // Clamp the source to what actually has inputs (unless the user is mid-op).
            const bool tOk = tC > 0, zOk = zC > 0;
            if (!s_operation_pending) {
                if (s_consolidate && zOk)      s_src = 1;               // bloat nudge → shielded
                else if (s_src == 0 && !tOk)   s_src = zOk ? 1 : 2;
                else if (s_src == 1 && !zOk)   s_src = tOk ? 0 : 2;
                else if (!tOk && zOk)          s_src = 1;
                else if (tOk && !zOk)          s_src = 0;
            }
        };
    });
}

void ShieldDialog::show(Mode mode)
{
    s_mode = mode;
    s_open = true;
    s_consolidate = false;           // reset preset flags so stale statics don't leak across opens
    s_src = 2;
    resetTransient();
    s_from_address[0] = '\0';
    if (mode == Mode::ShieldCoinbase) strncpy(s_from_address, "*", sizeof(s_from_address));
    s_to_address[0] = '\0';
    s_selected_zaddr_idx = -1;
    s_fee = DRAGONX_DEFAULT_FEE;
    s_utxo_limit = (int)schema::UI().drawElement("business", "utxo-limit").size;
    if (s_utxo_limit < 1) s_utxo_limit = 50;
    s_advanced = false;
    s_scope_loaded = false; s_scope_loading = false;
    s_t_count = s_z_count = 0; s_t_amount = s_z_amount = 0.0;
    s_last_poll = 0.0;
}

void ShieldDialog::showShieldCoinbase(const std::string& fromAddress)
{
    show(Mode::ShieldCoinbase);
    strncpy(s_from_address, fromAddress.c_str(), sizeof(s_from_address) - 1);
}

void ShieldDialog::showMerge()
{
    show(Mode::MergeToAddress);
    s_consolidate = false;
    s_src = 2;                       // generic merge: both sources
}

void ShieldDialog::showConsolidate()
{
    show(Mode::MergeToAddress);
    s_consolidate = true;
    s_src = 1;                       // wallet-bloat consolidation targets shielded notes (witness bloat)
}

void ShieldDialog::hide()
{
    s_open = false;
    resetTransient();
}

// Relevant count/amount for the currently-selected merge source.
static int    srcCount()  { return s_src == 0 ? s_t_count : s_src == 1 ? s_z_count : (s_t_count + s_z_count); }
static double srcAmount() { return s_src == 0 ? s_t_amount : s_src == 1 ? s_z_amount : (s_t_amount + s_z_amount); }

static std::string fmtAmt(double v) { char b[48]; std::snprintf(b, sizeof(b), "%.4f", v); return b; }

static std::string shortAddr(const std::string& a)
{
    if (a.size() <= 20) return a;
    return a.substr(0, 10) + "…" + a.substr(a.size() - 8);
}

// Auto-pick the best spendable z-address as the default destination (fewest hops for the user).
static void autoSelectDestination(const WalletState& state)
{
    if (s_to_address[0] != '\0' || state.z_addresses.empty()) return;
    int idx = bestSpendableAddressIndex(state.z_addresses);
    if (idx < 0) idx = 0;
    s_selected_zaddr_idx = idx;
    strncpy(s_to_address, state.z_addresses[idx].address.c_str(), sizeof(s_to_address) - 1);
}

// Fire the actual shield/merge op. Registers the opid with the shared poller (for balance refresh)
// AND kicks the modal's own live-progress poll.
static void submitOperation(App* app)
{
    s_operation_pending = true;
    s_op_terminal = false;
    s_status_message = TR("shield_submitting");
    s_last_poll = ImGui::GetTime();

    if (s_mode == ShieldDialog::Mode::ShieldCoinbase) {
        std::string from(s_from_address), to(s_to_address);
        double fee = s_fee; int limit = s_utxo_limit;
        if (!app->worker()) return;
        app->worker()->post([app, rpc = app->rpc(), from, to, fee, limit]() -> rpc::RPCWorker::MainCb {
            nlohmann::json result; std::string error;
            try {
                rpc::RPCClient::TraceScope trace("Shield dialog / Shield coinbase");
                result = rpc->call("z_shieldcoinbase", {from, to, fee, limit});
            } catch (const std::exception& e) { error = e.what(); }
            return [app, result, error]() {
                if (error.empty()) {
                    s_operation_id = result.value("opid", "");
                    s_status_message = TR("merge_progress");
                    Notifications::instance().success(TR("shield_started"));
                    app->trackOperation(s_operation_id);
                } else {
                    s_operation_pending = false; s_op_terminal = true;
                    s_status_message = std::string(TR("shield_error_prefix")) + error;
                    Notifications::instance().error(std::string(TR("shield_send_failed")) + error);
                }
            };
        });
        return;
    }

    // Merge / consolidate. Source → z_mergetoaddress fromaddress selector (this is the fix: shielded
    // notes, not just transparent UTXOs — the wallet-bloat the nudge warns about is shielded witnesses).
    std::vector<std::string> fromAddrs;
    if      (s_src == 0) fromAddrs = { "ANY_TADDR" };
    else if (s_src == 1) fromAddrs = { "ANY_SAPLING" };
    else                 fromAddrs = { "*" };
    std::string to(s_to_address);
    double fee = s_fee; int limit = s_utxo_limit;
    if (!app->worker()) return;
    app->worker()->post([app, rpc = app->rpc(), fromAddrs, to, fee, limit]() -> rpc::RPCWorker::MainCb {
        nlohmann::json addrs = nlohmann::json::array();
        for (const auto& a : fromAddrs) addrs.push_back(a);
        nlohmann::json result; std::string error;
        try {
            rpc::RPCClient::TraceScope trace("Shield dialog / Consolidate");
            // fromaddrs, toaddr, fee, transparent_limit, shielded_limit — cap both to the batch size.
            result = rpc->call("z_mergetoaddress", {addrs, to, fee, limit, limit});
        } catch (const std::exception& e) { error = e.what(); }
        return [app, result, error]() {
            if (error.empty()) {
                s_operation_id = result.value("opid", "");
                s_status_message = TR("merge_progress");
                Notifications::instance().success(TR("merge_started"));
                app->trackOperation(s_operation_id);
            } else {
                s_operation_pending = false; s_op_terminal = true;
                s_status_message = std::string(TR("shield_error_prefix")) + error;
                Notifications::instance().error(std::string(TR("merge_send_failed")) + error);
            }
        };
    });
}

// Live-progress self-poll: while an op is in flight, poll z_getoperationstatus every ~2s so the modal
// shows "Consolidating… → Done/Failed" without a manual button. (The shared poller also tracks it for
// balance refresh; this drives only the inline display.)
static void pollOperation(App* app)
{
    if (s_operation_id.empty() || s_op_terminal || !app->worker()) return;
    const double now = ImGui::GetTime();
    if (now - s_last_poll < 2.0) return;
    s_last_poll = now;
    std::string opid = s_operation_id;
    app->worker()->post([rpc = app->rpc(), opid]() -> rpc::RPCWorker::MainCb {
        nlohmann::json result; std::string error;
        try {
            rpc::RPCClient::TraceScope trace("Shield dialog / Op status");
            result = rpc->call("z_getoperationstatus", {nlohmann::json::array({opid})});
        } catch (const std::exception& e) { error = e.what(); }
        return [result, error]() {
            if (!error.empty() || !result.is_array() || result.empty()) return;   // transient — retry next tick
            const auto& op = result[0];
            const std::string status = op.value("status", "");
            if (status == "success") {
                s_operation_pending = false; s_op_terminal = true;
                s_status_message = TR("shield_completed");
                Notifications::instance().success(TR("shield_merge_done"));
            } else if (status == "failed") {
                std::string msg = op.value("error", nlohmann::json{}).value("message", std::string(TR("shield_unknown_error")));
                s_operation_pending = false; s_op_terminal = true;
                s_status_message = std::string(TR("shield_op_failed")) + msg;
                Notifications::instance().error(std::string(TR("shield_op_failed")) + msg);
            }
            // queued / executing → leave the "Consolidating…" message and keep polling.
        };
    });
}

void ShieldDialog::render(App* app)
{
    if (!s_open) return;

    auto& S = schema::UI();
    auto win        = S.window("dialogs.shield");
    auto addrLbl    = S.label("dialogs.shield", "address-label");
    auto addrFront  = S.label("dialogs.shield", "address-front-label");
    auto addrBack   = S.label("dialogs.shield", "address-back-label");
    auto feeInput   = S.input("dialogs.shield", "fee-input");
    auto utxoInput  = S.input("dialogs.shield", "utxo-limit-input");
    auto shieldBtn  = S.button("dialogs.shield", "shield-button");
    auto cancelBtn  = S.button("dialogs.shield", "cancel-button");
    const float dp  = Layout::dpiScale();
    const bool  isMerge = (s_mode == Mode::MergeToAddress);

    const char* title = s_consolidate ? TR("consolidate_title")
                      : isMerge        ? TR("merge_title")
                                       : TR("shield_title");

    material::OverlayDialogSpec ov;
    ov.title = title; ov.p_open = &s_open;
    ov.style = material::OverlayStyle::BlurFloat;
    ov.cardWidth = win.width; ov.idSuffix = "shielddialog";
    if (!material::BeginOverlayDialog(ov)) return;

    const auto& state = app->getWalletState();
    autoSelectDestination(state);
    pollOperation(app);
    if (isMerge && !s_scope_loaded && !s_scope_loading && s_operation_id.empty()) loadScope(app);

    // ── Description ──────────────────────────────────────────────────────────────────────────────
    ImGui::TextWrapped("%s", s_consolidate ? TR("consolidate_desc")
                            : isMerge       ? TR("merge_description")
                                            : TR("shield_description"));
    ImGui::Spacing();

    const bool opInFlight = !s_operation_id.empty();   // submitted — inputs frozen, showing progress

    // ── Merge: scope + source selector ──────────────────────────────────────────────────────────
    if (isMerge && !opInFlight) {
        if (s_scope_loading) {
            material::Type().textColored(material::TypeStyle::Caption, material::OnSurfaceMedium(),
                                         TR("merge_scope_loading"));
        } else if (s_scope_loaded) {
            char buf[160];
            std::snprintf(buf, sizeof(buf), TR("merge_scope_fmt"),
                          s_t_count, s_z_count, fmtAmt(s_t_amount + s_z_amount).c_str());
            material::Type().textColored(material::TypeStyle::Caption, material::OnSurfaceMedium(), buf);
        }
        ImGui::Spacing();

        // Source selector — only offer the types that actually have inputs.
        const bool tOk = s_t_count > 0, zOk = s_z_count > 0;
        if (tOk && zOk) {
            ImGui::AlignTextToFramePadding();
            ImGui::TextUnformatted(TR("merge_source"));
            ImGui::SameLine(0, Layout::spacingLg());
            ImGui::RadioButton(TR("merge_src_shielded"), &s_src, 1); ImGui::SameLine();
            ImGui::RadioButton(TR("merge_src_transparent"), &s_src, 0); ImGui::SameLine();
            ImGui::RadioButton(TR("merge_src_both"), &s_src, 2);
            ImGui::Spacing();
        }
    }

    // ── Shield coinbase: from address ───────────────────────────────────────────────────────────
    if (!isMerge && !opInFlight) {
        material::LabeledInput(TR("shield_from_address"), "##FromAddr", s_from_address, sizeof(s_from_address));
        ImGui::TextDisabled("%s", TR("shield_wildcard_hint"));
        ImGui::Spacing();
    }

    // ── Destination (z-address) ─────────────────────────────────────────────────────────────────
    if (!opInFlight) {
        ImGui::TextUnformatted(TR("shield_to_address"));
        if (state.z_addresses.empty()) {
            material::Type().textColored(material::TypeStyle::Caption, material::Warning(), TR("shield_no_zaddr_hint"));
            ImGui::Spacing();
            if (s_creating_addr) {
                ImGui::TextDisabled("%s", TR("merge_creating"));
            } else if (material::TactileButton(TR("merge_create_zaddr"), ImVec2(0, 0), S.resolveFont(shieldBtn.font))) {
                s_creating_addr = true;
                if (app->worker()) app->worker()->post([app, rpc = app->rpc()]() -> rpc::RPCWorker::MainCb {
                    std::string addr, error;
                    try {
                        rpc::RPCClient::TraceScope trace("Shield dialog / New z-address");
                        addr = rpc->call("z_getnewaddress", nlohmann::json::array()).get<std::string>();
                    } catch (const std::exception& e) { error = e.what(); }
                    return [app, addr, error]() {
                        s_creating_addr = false;
                        if (error.empty() && !addr.empty()) {
                            strncpy(s_to_address, addr.c_str(), sizeof(s_to_address) - 1);
                            Notifications::instance().success(TR("merge_addr_created"));
                        } else {
                            Notifications::instance().error(std::string(TR("shield_error_prefix")) + error);
                        }
                    };
                });
            }
        } else {
            std::string to_display = s_to_address[0] ? s_to_address : TR("shield_select_z");
            if (to_display.length() > static_cast<size_t>(addrLbl.truncate))
                to_display = to_display.substr(0, addrFront.truncate) + "..." + to_display.substr(to_display.length() - addrBack.truncate);
            ImGui::SetNextItemWidth(-1);
            if (ImGui::BeginCombo("##ToAddr", to_display.c_str())) {
                for (size_t i = 0; i < state.z_addresses.size(); i++) {
                    std::string label = state.z_addresses[i].address;
                    if (label.length() > static_cast<size_t>(addrLbl.truncate))
                        label = label.substr(0, addrFront.truncate) + "..." + label.substr(label.length() - addrBack.truncate);
                    bool selected = (s_selected_zaddr_idx == static_cast<int>(i));
                    if (ImGui::Selectable(label.c_str(), selected)) {
                        s_selected_zaddr_idx = static_cast<int>(i);
                        strncpy(s_to_address, state.z_addresses[i].address.c_str(), sizeof(s_to_address) - 1);
                    }
                    if (selected) ImGui::SetItemDefaultFocus();
                }
                ImGui::EndCombo();
            }
        }
        ImGui::Spacing();

        // ── Advanced (fee + batch size) ─────────────────────────────────────────────────────────
        ImDrawList* dl = ImGui::GetWindowDrawList();
        material::CollapsibleHeader(dl, "##AdvToggle", TR("merge_advanced"), s_advanced,
                                    ImGui::GetContentRegionAvail().x, material::Type().caption(),
                                    material::OnSurfaceMedium());
        if (s_advanced) {
            ImGui::Spacing();
            ImGui::TextUnformatted(TR("fee_label"));
            ImGui::SetNextItemWidth(feeInput.width * dp);
            ImGui::InputDouble("##Fee", &s_fee, 0.0001, 0.001, "%.8f");
            if (s_fee < 0.0) s_fee = 0.0;
            if (s_fee > 1.0) s_fee = 1.0;
            ImGui::SameLine(); ImGui::TextDisabled("DRGX");
            material::Type().textColored(material::TypeStyle::Caption, material::OnSurfaceMedium(), TR("merge_fee_hint"));

            ImGui::Spacing();
            ImGui::TextUnformatted(TR("merge_max_inputs"));
            ImGui::SetNextItemWidth(utxoInput.width * dp);
            ImGui::InputInt("##Limit", &s_utxo_limit);
            if (s_utxo_limit < 1)   s_utxo_limit = 1;
            if (s_utxo_limit > 100) s_utxo_limit = 100;
        }

        // Batch hint: one run only merges up to the limit; large sets need repeats.
        if (isMerge && s_scope_loaded && srcCount() > s_utxo_limit) {
            char hb[160];
            std::snprintf(hb, sizeof(hb), TR("merge_batch_fmt"), s_utxo_limit);
            ImGui::Spacing();
            material::Type().textColored(material::TypeStyle::Caption, material::Warning(), hb);
        }
        ImGui::Spacing();
    }

    // ── Live progress / status ──────────────────────────────────────────────────────────────────
    if (!s_status_message.empty()) {
        if (s_operation_pending && !s_op_terminal)
            ImGui::TextColored(ImVec4(1.0f, 0.8f, 0.0f, 1.0f), "%s", s_status_message.c_str());
        else
            ImGui::TextWrapped("%s", s_status_message.c_str());
        ImGui::Spacing();
    }

    // ── Footer ──────────────────────────────────────────────────────────────────────────────────
    const bool connected = app->isConnected();
    const bool syncing    = state.sync.syncing;
    const bool haveDest   = s_to_address[0] != '\0';

    if (opInFlight) {
        // After submit: just a Close button (progress shows above; op continues in the background).
        material::BeginOverlayDialogFooter(cancelBtn.width, /*drawSeparator=*/false);
        if (material::TactileButton(s_op_terminal ? TR("done") : TR("close"),
                                    ImVec2(cancelBtn.width, 0), S.resolveFont(cancelBtn.font)))
            s_open = false;
        material::EndOverlayDialog();
        return;
    }

    const char* primaryLabel = s_confirm ? TR("merge_confirm_btn")
                             : s_consolidate ? TR("consolidate_funds_btn")
                             : isMerge ? TR("merge_funds")
                                       : TR("shield_funds");
    const char* secondaryLabel = s_confirm ? TR("merge_back") : TR("cancel");

    // Confirm summary (inline, before the fund-moving call). Merge/consolidate shows amount + input
    // count; shield-coinbase just gets the button relabel (its inputs aren't enumerated here).
    if (s_confirm && isMerge) {
        char cb[200];
        std::snprintf(cb, sizeof(cb), TR("merge_confirm_fmt"),
                      fmtAmt(srcAmount()).c_str(), srcCount(), shortAddr(s_to_address).c_str());
        ImGui::TextWrapped("%s", cb);
        ImGui::Spacing();
    }

    bool can_submit = haveDest && connected && !syncing;
    if (isMerge && s_scope_loaded && srcCount() == 0) can_submit = false;

    float footerBtnW = shieldBtn.width + cancelBtn.width + ImGui::GetStyle().ItemSpacing.x;
    material::BeginOverlayDialogFooter(footerBtnW, /*drawSeparator=*/false);

    if (!can_submit) ImGui::BeginDisabled();
    if (material::TactileButton(primaryLabel, ImVec2(shieldBtn.width, 0), S.resolveFont(shieldBtn.font))) {
        if (s_confirm) { submitOperation(app); }
        else           { s_confirm = true; }     // first click → show the confirm summary
    }
    if (!can_submit) ImGui::EndDisabled();
    if (!can_submit && ImGui::IsItemHovered(ImGuiHoveredFlags_AllowWhenDisabled)) {
        if (!connected)                              material::Tooltip("%s", TR("send_tooltip_not_connected"));
        else if (syncing)                            material::Tooltip("%s", TR("send_tooltip_syncing"));
        else if (!haveDest)                          material::Tooltip("%s", TR("shield_select_z"));
        else if (isMerge && srcCount() == 0)         material::Tooltip("%s", TR("merge_no_spendable"));
    }

    ImGui::SameLine();
    if (material::TactileButton(secondaryLabel, ImVec2(cancelBtn.width, 0), S.resolveFont(cancelBtn.font))) {
        if (s_confirm) s_confirm = false;         // Back → return to the form
        else           s_open = false;            // Cancel → close
    }

    material::EndOverlayDialog();
}

} // namespace ui
} // namespace dragonx
