// DragonX Wallet - ImGui Edition
// Copyright 2024-2026 The Hush Developers
// Released under the GPLv3
//
// wallet_file_probe.h — read light metadata off a wallet.dat WITHOUT loading it into the daemon.
//
// DragonX wallet.dat is a Berkeley DB (btree) key/value store. Only the private-key material is
// encrypted (ckey/czkey/csapzkey/chdseed); the record *keys* and non-secret metadata live in
// plaintext, so a handful of boolean facts can be recovered cheaply by (a) validating the BDB btree
// magic and (b) scanning the raw bytes for the length-prefixed record names the daemon writes
// (a std::string is serialized as CompactSize(len)+bytes, so e.g. an "mkey" record's key begins with
// the 5 bytes 0x04 'm' 'k' 'e' 'y'). This is a heuristic PRESENCE test — reliable for booleans, not
// for exact counts — with no libdb dependency and no daemon. It never reads or exposes key material.

#pragma once

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <fstream>
#include <string>
#include <utility>
#include <vector>

namespace dragonx {
namespace util {

struct WalletFileProbe {
    bool isBerkeleyDB = false;  ///< file has a valid BDB btree metapage magic (looks like a real wallet.dat)
    bool encrypted    = false;  ///< has an "mkey" master-key record → passphrase-encrypted
    bool hdSeed       = false;  ///< has hdseed/chdseed/hdchain → HD/seed wallet (else legacy)
    bool hasShielded  = false;  ///< has zkey/czkey/sapzkey/csapzkey → holds shielded addresses
    // A found marker is always definitive; a MISSING marker is only trustworthy when the scan actually
    // covered the whole file. `scanComplete` is true when we reached EOF, or early-exited having found
    // everything — false only when the byte cap cut the scan short. Callers should treat negative flags
    // (e.g. "legacy" = !hdSeed, or "not encrypted" = !encrypted) as reliable only when scanComplete is true.
    bool scanComplete = false;
    std::size_t bytesRead = 0;  ///< bytes actually consumed (early-exit reads far less than the file) — for budgeting
};

// Probe a wallet.dat by header-validating it as a BDB btree then byte-scanning (bounded, streaming,
// early-exit) for record markers. `maxBytes` caps how far we read so a pathologically large file can't
// stall the caller. Returns an all-false probe (isBerkeleyDB=false) for anything that isn't a readable
// BDB btree file. Safe to call on a wallet currently open by the daemon (read-only, pattern scan only).
inline WalletFileProbe probeWalletFile(const std::string& path,
                                       std::size_t maxBytes = 96u * 1024u * 1024u) {
    WalletFileProbe out;
    std::ifstream f(path, std::ios::binary);
    if (!f) return out;

    // --- 1) Validate the Berkeley DB btree metapage magic (offset 12, either endianness). ---
    unsigned char hdr[16];
    f.read(reinterpret_cast<char*>(hdr), sizeof(hdr));
    if (f.gcount() < static_cast<std::streamsize>(sizeof(hdr))) return out;
    auto rd32 = [](const unsigned char* p, bool le) -> uint32_t {
        return le ? (uint32_t)p[0] | ((uint32_t)p[1] << 8) | ((uint32_t)p[2] << 16) | ((uint32_t)p[3] << 24)
                  : (uint32_t)p[3] | ((uint32_t)p[2] << 8) | ((uint32_t)p[1] << 16) | ((uint32_t)p[0] << 24);
    };
    constexpr uint32_t kBtreeMagic = 0x00053162u;  // DB_BTREEMAGIC
    if (rd32(hdr + 12, true) != kBtreeMagic && rd32(hdr + 12, false) != kBtreeMagic) return out;
    out.isBerkeleyDB = true;

    // --- 2) Streaming byte-scan for length-prefixed record markers. ---
    // group: 0 = encrypted, 1 = hd/seed, 2 = shielded.
    const std::vector<std::pair<std::string, int>> patterns = {
        {std::string("\x04", 1) + "mkey",     0},  // master key → wallet is encrypted
        {std::string("\x06", 1) + "hdseed",   1},  // plaintext HD seed
        {std::string("\x07", 1) + "chdseed",  1},  // encrypted HD seed
        {std::string("\x07", 1) + "hdchain",  1},  // HD chain counter (present in both HD variants)
        {std::string("\x04", 1) + "zkey",     2},  // sprout shielded key
        {std::string("\x05", 1) + "czkey",    2},  // encrypted sprout key
        {std::string("\x07", 1) + "sapzkey",  2},  // sapling shielded key
        {std::string("\x08", 1) + "csapzkey", 2},  // encrypted sapling key
    };
    std::size_t maxPat = 0;
    for (const auto& p : patterns) maxPat = std::max(maxPat, p.first.size());
    const std::size_t overlap = maxPat > 0 ? maxPat - 1 : 0;  // bytes carried between chunks for split matches

    f.clear();
    f.seekg(0, std::ios::beg);
    constexpr std::size_t kChunk = 1u << 20;  // 1 MiB
    std::vector<char> buf(kChunk);
    std::string carry;
    std::size_t readTotal = 0;
    bool enc = false, hd = false, sh = false;
    bool eof = false, allFound = false;
    while (readTotal < maxBytes) {
        const std::size_t want = std::min(kChunk, maxBytes - readTotal);
        f.read(buf.data(), static_cast<std::streamsize>(want));
        const std::streamsize got = f.gcount();
        if (got <= 0) { eof = true; break; }
        readTotal += static_cast<std::size_t>(got);

        std::string window;
        window.reserve(carry.size() + static_cast<std::size_t>(got));
        window.assign(carry);
        window.append(buf.data(), static_cast<std::size_t>(got));

        for (const auto& p : patterns) {
            bool& flag = (p.second == 0) ? enc : (p.second == 1) ? hd : sh;
            if (!flag && window.find(p.first) != std::string::npos) flag = true;
        }
        if (enc && hd && sh) { allFound = true; break; }  // nothing more to learn

        if (window.size() > overlap) carry.assign(window.data() + window.size() - overlap, overlap);
        else                         carry.assign(window);
        if (static_cast<std::size_t>(got) < want) { eof = true; break; }  // reached EOF
    }
    out.encrypted    = enc;
    out.hdSeed       = hd;
    out.hasShielded  = sh;
    out.bytesRead    = readTotal;
    // false only if the byte cap cut a larger file short. If we stopped because readTotal hit maxBytes on
    // an exact chunk boundary, ifstream never set eofbit — peek() tells us whether the file actually ended
    // there (EOF → complete) or there's more beyond the cap (a real byte → truncated).
    out.scanComplete = allFound || eof;
    if (!out.scanComplete && f && f.peek() == std::ifstream::traits_type::eof()) out.scanComplete = true;
    return out;
}

// ─────────────────────────────────────────────────────────────────────────────────────────────────
// Tier 2: exact record counts by actually walking the Berkeley DB btree (still no daemon, no libdb).
//
// wallet.dat is a BDB btree: a metapage (page 0) points at a root page; internal pages point at child
// pages; leaf pages hold (key,data) item pairs. We traverse from the root, visiting only reachable
// leaves (so freed/stale pages aren't counted), and tally records by the length-prefixed type name each
// key begins with. Every offset is bounds-checked against the page/file; a visited-set + page/key caps
// make it safe on a corrupt or adversarial file. On any structural surprise it returns parsed=false and
// the caller falls back to the byte-scan probe above.
struct WalletBtreeStats {
    bool parsed          = false;  ///< the btree walked cleanly
    bool complete        = false;  ///< the whole file was read (not cap-truncated) → counts are exact
    int  transparentKeys = 0;      ///< key + wkey + ckey  (spendable keys, incl. change; keypool "pool" recs excluded)
    int  shieldedKeys    = 0;      ///< zkey + czkey + sapzkey + csapzkey  (≈ shielded addresses)
    int  addressBook     = 0;      ///< name records (labeled/received addresses)
    int  txCount         = 0;      ///< tx records (wallet transactions)
    long long createdEpoch = 0;    ///< earliest keymeta nCreateTime (wallet birthday); 0 = unknown
    bool encrypted       = false;  ///< saw an mkey record
    bool hdSeed          = false;  ///< saw hdseed/chdseed/hdchain
    // BIP39 seed-phrase detection read straight off the hdchain record's fMnemonicSeed flag (which the
    // daemon keeps PLAINTEXT even in an encrypted wallet — only the seed itself is crypted):
    // 0 = undecidable, 1 = mnemonic seed-phrase wallet (z_exportmnemonic works), 2 = HD/legacy, no phrase.
    int  mnemonicSeed    = 0;
    std::size_t bytesRead = 0;     ///< bytes actually read (for budgeting)
    int  addresses() const { return transparentKeys + shieldedKeys; }
};

// Recover the BIP39-mnemonic flag from a serialized CHDChain record value (the plaintext "hdchain"
// record). The daemon serializes CHDChain in Bitcoin/Zcash little-endian byte order as:
//   nVersion:int32 | seedFp:32B | nCreateTime:int64 | saplingAccountCounter:uint32
//   | [nVersion>=2 (VERSION_HD_TRANSPARENT)] transparentChildCounter:uint32
//   | [nVersion>=3 (VERSION_HD_MNEMONIC)]    fMnemonicSeed:bool(1B)
// so fMnemonicSeed sits at byte offset 52 (= 4+32+8+4+4) once nVersion>=3. Returns 1 = mnemonic,
// 2 = HD but not mnemonic (raw-entropy seed, or a pre-mnemonic v1/v2 wallet), 0 = undecidable.
inline int hdChainMnemonicFlag(const unsigned char* val, uint32_t len) {
    if (!val || len < 4) return 0;
    const uint32_t v = (uint32_t)val[0] | ((uint32_t)val[1] << 8) |
                       ((uint32_t)val[2] << 16) | ((uint32_t)val[3] << 24);
    const int32_t nVersion = (int32_t)v;
    if (nVersion < 1 || nVersion > 100) return 0;   // implausible version → treat as unreadable, not "legacy"
    if (nVersion < 3) return 2;                      // VERSION_HD_MNEMONIC not reached → field absent → no phrase
    constexpr uint32_t kMnemonicOff = 52;
    if (len <= kMnemonicOff) return 0;              // truncated/unexpected value → can't read the flag
    return val[kMnemonicOff] ? 1 : 2;
}

inline WalletBtreeStats parseWalletBtree(const std::string& path,
                                         std::size_t maxBytes = 256u * 1024u * 1024u) {
    WalletBtreeStats st;
    std::ifstream f(path, std::ios::binary);
    if (!f) return st;
    std::string buf;
    {
        f.seekg(0, std::ios::end);
        std::streamoff sz = f.tellg();
        if (sz < 512) return st;
        const std::size_t want = std::min<std::size_t>(static_cast<std::size_t>(sz), maxBytes);
        f.seekg(0, std::ios::beg);
        buf.resize(want);
        f.read(&buf[0], static_cast<std::streamsize>(want));
        buf.resize(static_cast<std::size_t>(std::max<std::streamsize>(0, f.gcount())));
        if (buf.size() < 512) return st;
        st.bytesRead = buf.size();
        st.complete  = (buf.size() == static_cast<std::size_t>(sz));   // read the whole file, not cap-truncated
    }
    const unsigned char* B = reinterpret_cast<const unsigned char*>(buf.data());
    const std::size_t N = buf.size();

    // Byte order comes from the metapage magic (BDB stores fields in the creating machine's order).
    auto rd32at = [&](std::size_t o, bool le) -> uint32_t {
        return le ? (uint32_t)B[o] | ((uint32_t)B[o+1]<<8) | ((uint32_t)B[o+2]<<16) | ((uint32_t)B[o+3]<<24)
                  : (uint32_t)B[o+3] | ((uint32_t)B[o+2]<<8) | ((uint32_t)B[o+1]<<16) | ((uint32_t)B[o]<<24);
    };
    constexpr uint32_t kBtreeMagic = 0x00053162u;
    bool le;
    if      (rd32at(12, true)  == kBtreeMagic) le = true;
    else if (rd32at(12, false) == kBtreeMagic) le = false;
    else return st;   // not a BDB btree

    auto r32 = [&](std::size_t o) { return o + 4 <= N ? rd32at(o, le) : 0u; };
    auto r16 = [&](std::size_t o) -> uint32_t {
        if (o + 2 > N) return 0;
        return le ? (uint32_t)B[o] | ((uint32_t)B[o+1]<<8) : (uint32_t)B[o+1] | ((uint32_t)B[o]<<8);
    };

    const uint32_t pagesize = r32(20);
    if (pagesize < 512 || pagesize > 65536 || (pagesize & (pagesize - 1)) != 0) return st;  // must be a power of two
    const uint32_t npages   = static_cast<uint32_t>(N / pagesize);
    const uint32_t root     = r32(88);
    if (npages == 0 || root == 0 || root >= npages) return st;
    // Page-level checksums (DB_CHKSUM) or encryption (DB_ENCRYPT) shift the per-page offset index array
    // (26 → 32 / 64), which our fixed-offset-26 reads don't account for — that would silently under-count
    // rather than fail. DragonX wallets use neither, so bail to the byte-scan fallback if either is set
    // (metapage encrypt_alg@24, metaflags@26 & DBMETA_CHKSUM 0x01).
    if (B[24] != 0 || (B[26] & 0x01)) return st;

    // BDB page layout: hdr[26] = {..., entries@20:u16, level@24:u8, type@25:u8}, then a u16 offset array.
    // Items on a leaf are BKEYDATA {len@0:u16, type@2:u8, data@3}; on an internal page BINTERNAL {pgno@4}.
    constexpr uint8_t P_IBTREE = 3, P_LBTREE = 5, P_BTREEMETA = 9, B_KEYDATA = 1;
    constexpr std::size_t kMaxPagesVisited = 600000;   // bounds a corrupt/huge file
    constexpr int         kMaxKeys         = 4000000;

    std::vector<bool> visited(npages, false);
    std::size_t pagesVisited = 0;
    int keys = 0;
    bool aborted = false;
    auto rdpgno = [&](const unsigned char* p) -> uint32_t {
        return le ? (uint32_t)p[0] | ((uint32_t)p[1]<<8) | ((uint32_t)p[2]<<16) | ((uint32_t)p[3]<<24)
                  : (uint32_t)p[3] | ((uint32_t)p[2]<<8) | ((uint32_t)p[1]<<16) | ((uint32_t)p[0]<<24);
    };

    // Walk the btree rooted at `rootPg`, invoking fn(keyPtr,keyLen, dataPtr,dataLen,dataType) per leaf pair.
    auto traverse = [&](uint32_t rootPg, auto&& fn) {
        std::fill(visited.begin(), visited.end(), false);
        std::vector<uint32_t> stack;
        // Mark pages visited at PUSH time (here + at each internal child below) so every page is enqueued
        // at most once. The stack then stays O(npages) — a crafted file with many internal pages all
        // referencing a shared child set can't balloon it to ~npages*entries pushes (a ~0.5 GB DoS).
        if (rootPg < npages && !visited[rootPg]) { visited[rootPg] = true; stack.push_back(rootPg); }
        while (!stack.empty()) {
            const uint32_t pg = stack.back(); stack.pop_back();
            if (pg >= npages) continue;
            if (++pagesVisited > kMaxPagesVisited) { aborted = true; return; }
            const std::size_t base = static_cast<std::size_t>(pg) * pagesize;
            if (base + 26 > N) continue;
            const uint8_t  type    = B[base + 25];
            const uint32_t entries = r16(base + 20);
            if (26 + static_cast<std::size_t>(entries) * 2 > pagesize) continue;   // index array must fit
            if (type == P_IBTREE) {
                for (uint32_t i = 0; i < entries; ++i) {
                    const uint32_t off = r16(base + 26 + i * 2);
                    if (off + 8 > pagesize) continue;
                    const uint32_t child = r32(base + off + 4);
                    if (child > 0 && child < npages && !visited[child]) { visited[child] = true; stack.push_back(child); }
                }
            } else if (type == P_LBTREE) {
                for (uint32_t i = 0; i + 1 < entries; i += 2) {   // items alternate (key, data)
                    if (++keys > kMaxKeys) { aborted = true; return; }
                    const uint32_t ko = r16(base + 26 + i * 2);
                    const uint32_t dO = r16(base + 26 + (i + 1) * 2);
                    if (ko + 3 > pagesize || dO + 3 > pagesize) continue;
                    if (B[base + ko + 2] != B_KEYDATA) continue;   // overflow/dup key — never a type key
                    const uint32_t kl = r16(base + ko);
                    if (kl < 1 || ko + 3 + kl > pagesize) continue;
                    const uint8_t  dtype = B[base + dO + 2];
                    const uint32_t dl    = r16(base + dO);
                    const unsigned char* dp = (dO + 3 + dl <= pagesize) ? B + base + dO + 3 : nullptr;
                    fn(B + base + ko + 3, kl, dp, dl, dtype);
                }
            }
            // metapage / overflow / free pages: ignored.
        }
    };

    // wallet.dat stores its records in a NAMED sub-database ("main"): the file's root btree is a master map
    // of subdb-name -> subdb meta/root pgno. Follow each mapping to the real record btree. (A plain,
    // single-database BDB file has no such mapping, so we fall back to walking the file root directly.)
    std::vector<uint32_t> subRoots;
    traverse(root, [&](const unsigned char*, uint32_t, const unsigned char* dp, uint32_t dl, uint8_t dtype) {
        if (dtype != B_KEYDATA || dl != 4 || !dp) return;
        // The master-db value (the subdb's meta/root pgno) is stored BIG-ENDIAN regardless of the file's
        // native order. Try big-endian first, then native, and accept whichever lands on a real subdb page.
        const uint32_t cand[2] = {
            (uint32_t)dp[3] | ((uint32_t)dp[2]<<8) | ((uint32_t)dp[1]<<16) | ((uint32_t)dp[0]<<24),  // big-endian
            rdpgno(dp),                                                                              // native
        };
        for (const uint32_t pgno : cand) {
            if (pgno == 0 || pgno >= npages) continue;
            const uint8_t pt = B[static_cast<std::size_t>(pgno) * pagesize + 25];
            if (pt == P_BTREEMETA) {                              // subdb metapage → its root is at +88
                const uint32_t sr = r32(static_cast<std::size_t>(pgno) * pagesize + 88);
                if (sr > 0 && sr < npages) { subRoots.push_back(sr); break; }
            } else if (pt == P_LBTREE || pt == P_IBTREE) {        // mapping points straight at the root
                subRoots.push_back(pgno); break;
            }
        }
    });
    if (aborted) return st;
    if (subRoots.empty()) subRoots.push_back(root);              // single-database file

    auto countRecord = [&](const unsigned char* kp, uint32_t kl,
                           const unsigned char* dp, uint32_t dl, uint8_t dtype) {
        const uint32_t nlen = kp[0];                              // CompactSize length of the type string
        if (nlen < 2 || nlen > 20 || 1u + nlen > kl) return;
        const char* nm = reinterpret_cast<const char*>(kp + 1);
        auto is = [&](const char* s) { return std::strlen(s) == nlen && std::memcmp(nm, s, nlen) == 0; };
        if      (is("key") || is("wkey") || is("ckey"))                        st.transparentKeys++;
        else if (is("zkey") || is("czkey") || is("sapzkey") || is("csapzkey")) st.shieldedKeys++;
        else if (is("name"))                                                   st.addressBook++;
        else if (is("tx"))                                                     st.txCount++;
        else if (is("mkey"))                                                   st.encrypted = true;
        else if (is("hdseed") || is("chdseed"))                                st.hdSeed = true;
        else if (is("hdchain")) {
            st.hdSeed = true;
            // The hdchain VALUE carries fMnemonicSeed — read it to tell a BIP39 seed-phrase wallet
            // apart from a legacy/raw-entropy HD wallet (bare record presence can't). First one wins.
            if (st.mnemonicSeed == 0 && dtype == B_KEYDATA && dp)
                st.mnemonicSeed = hdChainMnemonicFlag(dp, dl);
        }
        else if (is("keymeta")) {
            // CKeyMetadata data = nVersion(int32) + nCreateTime(int64) + …, all little-endian (Bitcoin
            // serialization). The earliest non-zero nCreateTime is the wallet birthday (daemon's
            // nTimeFirstKey). B_KEYDATA(1) only; overflow values aren't inline.
            if (dtype == B_KEYDATA && dp && dl >= 12) {
                long long t = 0;
                for (int b = 0; b < 8; ++b) t |= static_cast<long long>(dp[4 + b]) << (8 * b);
                if (t > 0 && (st.createdEpoch == 0 || t < st.createdEpoch)) st.createdEpoch = t;
            }
        }
    };
    for (const uint32_t sr : subRoots) {
        traverse(sr, [&](const unsigned char* kp, uint32_t kl, const unsigned char* dp, uint32_t dl, uint8_t dt) {
            countRecord(kp, kl, dp, dl, dt);
        });
        if (aborted) return st;
    }
    st.parsed = true;
    return st;
}

}  // namespace util
}  // namespace dragonx