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

#include "export_all_keys_dialog.h"
#include "../../app.h"
#include <sodium.h>   // sodium_memzero — wipe exported key material from memory (B7)
#include "../../rpc/rpc_client.h"
#include "../../rpc/rpc_worker.h"
#include "../../util/i18n.h"
#include "../../util/platform.h"
#include "../notifications.h"
#include "../schema/ui_schema.h"
#include "../material/draw_helpers.h"
#include "../material/type.h"
#include "../widgets/copy_field.h"
#include "../theme.h"
#include "../../embedded/IconsMaterialDesign.h"
#include "imgui.h"

#include <string>
#include <vector>
#include <fstream>
#include <ctime>

namespace dragonx {
namespace ui {

using json = nlohmann::json;

// Static state
static bool s_open = false;
static bool s_exporting = false;
static std::string s_status;        // error / partial message (non-empty => red text)
static std::string s_saved_path;    // full path on full success (=> green check + copy field)
static std::string s_exported_keys;
static int s_total_addresses = 0;
static int s_exported_count = 0;
static bool s_include_z = true;
static bool s_include_t = true;
static char s_filename[256] = "";

void ExportAllKeysDialog::show()
{
    s_open = true;
    s_exporting = false;
    s_status.clear();
    s_saved_path.clear();
    s_exported_keys.clear();
    s_total_addresses = 0;
    s_exported_count = 0;
    s_include_z = true;
    s_include_t = true;
    
    // Generate default filename with timestamp
    std::time_t now = std::time(nullptr);
    char timebuf[32];
    std::strftime(timebuf, sizeof(timebuf), "%Y%m%d_%H%M%S", std::localtime(&now));
    snprintf(s_filename, sizeof(s_filename), "dragonx_keys_%s.txt", timebuf);
}

bool ExportAllKeysDialog::isOpen()
{
    return s_open;
}

void ExportAllKeysDialog::hide()
{
    s_open = false;
    s_status.clear();
    s_saved_path.clear();
}

void ExportAllKeysDialog::render(App* app)
{
    if (!s_open) return;
    namespace m = material;

    m::OverlayDialogSpec ov;
    ov.title     = TR("export_keys_title");
    ov.p_open    = &s_open;
    ov.style     = m::OverlayStyle::BlurFloat;
    ov.cardWidth = 600.0f;
    ov.idSuffix  = "exportallkeys";
    if (m::BeginOverlayDialog(ov)) {
        m::DialogWarningHeader(TR("export_keys_danger"));
        ImGui::Spacing();

        ImGui::BeginDisabled(s_exporting);
        ImGui::TextUnformatted(TR("export_keys_options"));
        ImGui::Checkbox(TR("export_keys_include_z"), &s_include_z);
        ImGui::Checkbox(TR("export_keys_include_t"), &s_include_t);
        ImGui::Spacing();
        m::LabeledInput(TR("output_filename"), "##Filename", s_filename, sizeof(s_filename));
        ImGui::Spacing();
        ImGui::TextColored(ImGui::ColorConvertU32ToFloat4(m::OnSurfaceDisabled()), "%s", TR("file_save_location"));
        ImGui::EndDisabled();

        // Progress / result.
        if (s_exporting) {
            ImGui::Spacing();
            ImGui::TextColored(ImGui::ColorConvertU32ToFloat4(m::Primary()), TR("export_keys_progress"),
                               s_exported_count, s_total_addresses);
            ImGui::SameLine(0, 0);
            ImGui::TextColored(ImGui::ColorConvertU32ToFloat4(m::Primary()), "%s", m::LoadingDots());
        } else if (!s_saved_path.empty()) {
            ImGui::Spacing();
            ImGui::PushFont(m::Type().iconSmall());
            ImGui::TextColored(m::SuccessVec4(), ICON_MD_CHECK_CIRCLE);
            ImGui::PopFont();
            ImGui::SameLine(0, 4.0f * Layout::dpiScale());
            ImGui::TextColored(m::SuccessVec4(), "%s", TR("export_keys_success"));
            widgets::AddressCopyField("##exportedkeyspath", s_saved_path);
        } else if (!s_status.empty()) {
            ImGui::Spacing();
            ImGui::PushTextWrapPos(0.0f);
            ImGui::TextColored(ImVec4(0.8f, 0.3f, 0.3f, 1.0f), "%s", s_status.c_str());
            ImGui::PopTextWrapPos();
        }

        ImGui::Spacing();

        bool doExport = false, doClose = false;
        m::DialogActionFooter(TR("export_keys_btn"), !s_exporting, TR("close"), doExport, doClose);
        if (doClose) s_open = false;
        if (doExport) {
            if (!s_include_z && !s_include_t) {
                Notifications::instance().warning(TR("export_keys_select_type"));
            } else if (!app->rpc() || !app->rpc()->isConnected()) {
                Notifications::instance().error(TR("export_keys_not_connected"));
            } else {
                s_exporting = true;
                s_exported_keys.clear();
                s_exported_count = 0;
                s_status.clear();
                s_saved_path.clear();

                const auto& state = app->getWalletState();
                
                // Count total addresses to export
                s_total_addresses = 0;
                if (s_include_z) s_total_addresses += static_cast<int>(state.z_addresses.size());
                if (s_include_t) s_total_addresses += static_cast<int>(state.t_addresses.size());
                
                if (s_total_addresses == 0) {
                    s_exporting = false;
                    s_status = TR("export_keys_none_addrs");
                    m::EndOverlayDialog();
                    return;
                }
                
                // Collect addresses to export (copy for worker thread)
                std::vector<std::string> z_addrs, t_addrs;
                if (s_include_z) {
                    for (const auto& a : state.z_addresses) z_addrs.push_back(a.address);
                }
                if (s_include_t) {
                    for (const auto& a : state.t_addresses) t_addrs.push_back(a.address);
                }
                std::string filename(s_filename);
                
                // Run all key exports on worker thread
                if (app->worker()) {
                    app->worker()->post([rpc = app->rpc(), z_addrs, t_addrs, filename]() -> rpc::RPCWorker::MainCb {
                        std::string keys;
                        int exported = 0;
                        int total = static_cast<int>(z_addrs.size() + t_addrs.size());
                        
                        // Header
                        keys = "# DragonX Wallet - Private Keys Export\n";
                        keys += "# Generated: ";
                        std::time_t now = std::time(nullptr);
                        char timebuf[64];
                        std::strftime(timebuf, sizeof(timebuf), "%Y-%m-%d %H:%M:%S\n", std::localtime(&now));
                        keys += timebuf;
                        keys += "# KEEP THIS FILE SECURE!\n\n";
                        
                        // Export Z-addresses
                        if (!z_addrs.empty()) {
                            keys += "# === Z-Addresses (Shielded) ===\n\n";
                            for (const auto& addr : z_addrs) {
                                try {
                                    rpc::RPCClient::TraceScope trace("Settings / Export all keys");
                                    std::string k = rpc->callSecretString("z_exportkey", {addr});  // scrubs raw body + json node (B7)
                                    keys += "# Address: " + addr + "\n";
                                    keys += k + "\n\n";
                                    exported++;   // count only real successes (locked/failed keys don't count)
                                    if (!k.empty()) sodium_memzero(&k[0], k.size());  // wipe the transient copy
                                } catch (...) {}  // non-string / locked key → skip (not counted)
                            }
                        }
                        
                        // Export T-addresses
                        if (!t_addrs.empty()) {
                            keys += "# === T-Addresses (Transparent) ===\n\n";
                            for (const auto& addr : t_addrs) {
                                try {
                                    rpc::RPCClient::TraceScope trace("Settings / Export all keys");
                                    std::string k = rpc->callSecretString("dumpprivkey", {addr});  // scrubs raw body + json node (B7)
                                    keys += "# Address: " + addr + "\n";
                                    keys += k + "\n\n";
                                    exported++;   // count only real successes (locked/failed keys don't count)
                                    if (!k.empty()) sodium_memzero(&k[0], k.size());  // wipe the transient copy
                                } catch (...) {}  // non-string / locked key → skip (not counted)
                            }
                        }
                        
                        // Only write a file if we actually exported keys — a keyless file would be a
                        // misleading "export" (the usual cause of 0 keys is an encrypted+locked wallet).
                        std::string configDir = util::Platform::getConfigDir();
                        std::string filepath = configDir + "/" + filename;
                        bool writeOk = false;
                        if (exported > 0) {
                            // Write the plaintext private keys 0600 + atomically (never
                            // world/group-readable, never a half-written file) — the same restricted
                            // atomic-write idiom the PIN vault uses. A default ofstream would create
                            // the key dump with umask-derived (often 0644) permissions.
                            writeOk = util::Platform::writeFileAtomically(filepath, keys, /*restrictPermissions=*/true);
                        }
                        if (!keys.empty()) sodium_memzero(&keys[0], keys.size());  // don't leave every key in freed heap

                        return [exported, total, filepath, writeOk]() {
                            s_exported_count = exported;
                            s_exporting = false;
                            char buf[320];
                            if (exported == 0) {
                                snprintf(buf, sizeof(buf), TR("export_keys_none_result"), total);
                                s_status = buf;
                                Notifications::instance().error(TR("export_keys_none_toast"));
                            } else if (!writeOk) {
                                s_status = TR("export_keys_write_fail");
                                Notifications::instance().error(TR("export_keys_write_fail"));
                            } else if (exported < total) {
                                snprintf(buf, sizeof(buf), TR("export_keys_partial"), exported, total);
                                s_status = buf;
                                snprintf(buf, sizeof(buf), TR("export_keys_partial_toast"), exported, total);
                                Notifications::instance().warning(buf);
                            } else {
                                s_saved_path = filepath;
                                Notifications::instance().success(TR("export_keys_success"), 5.0f);
                            }
                        };
                    });
                }
            }
        }

        m::EndOverlayDialog();
    }
}

} // namespace ui
} // namespace dragonx
