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

#include "address_validation.h"

#include <sodium.h>

#include <algorithm>
#include <cctype>
#include <cstdint>
#include <cstring>
#include <vector>

namespace dragonx {
namespace util {

namespace {

constexpr const char* kBase58 =
    "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz";
constexpr const char* kBech32 = "qpzry9x8gf2tvdw0s3jn54khce6mua7l";

// Decode a Base58 string to bytes (big-endian), preserving leading-zero bytes.
bool base58Decode(const std::string& s, std::vector<std::uint8_t>& out)
{
    std::vector<std::uint8_t> bytes;  // little-endian during accumulation
    for (char ch : s) {
        const char* p = std::strchr(kBase58, ch);
        if (p == nullptr || ch == '\0') return false;
        int carry = static_cast<int>(p - kBase58);
        for (auto& b : bytes) {
            carry += static_cast<int>(b) * 58;
            b = static_cast<std::uint8_t>(carry & 0xff);
            carry >>= 8;
        }
        while (carry > 0) {
            bytes.push_back(static_cast<std::uint8_t>(carry & 0xff));
            carry >>= 8;
        }
    }
    std::vector<std::uint8_t> result;
    for (char ch : s) {  // leading '1's map to leading zero bytes
        if (ch == '1') result.push_back(0);
        else break;
    }
    result.insert(result.end(), bytes.rbegin(), bytes.rend());
    out = std::move(result);
    return true;
}

std::uint32_t bech32Polymod(const std::vector<int>& values)
{
    static const std::uint32_t kGen[5] = {
        0x3b6a57b2, 0x26508e6d, 0x1ea119fa, 0x3d4233dd, 0x2a1462b3};
    std::uint32_t chk = 1;
    for (int v : values) {
        std::uint32_t top = chk >> 25;
        chk = ((chk & 0x1ffffff) << 5) ^ static_cast<std::uint32_t>(v);
        for (int i = 0; i < 5; ++i) {
            if ((top >> i) & 1) chk ^= kGen[i];
        }
    }
    return chk;
}

std::vector<int> bech32HrpExpand(const std::string& hrp)
{
    std::vector<int> out;
    out.reserve(hrp.size() * 2 + 1);
    for (char c : hrp) out.push_back(static_cast<unsigned char>(c) >> 5);
    out.push_back(0);
    for (char c : hrp) out.push_back(static_cast<unsigned char>(c) & 31);
    return out;
}

}  // namespace

bool decodeBase58Check(const std::string& s, std::vector<std::uint8_t>& payloadOut)
{
    if (s.size() < 5 || s.size() > 256) return false;
    std::vector<std::uint8_t> data;
    if (!base58Decode(s, data)) return false;
    if (data.size() < 5) return false;  // need at least version(1) + checksum(4)

    const std::size_t payloadLen = data.size() - 4;
    unsigned char h1[crypto_hash_sha256_BYTES];
    unsigned char h2[crypto_hash_sha256_BYTES];
    crypto_hash_sha256(h1, data.data(), payloadLen);
    crypto_hash_sha256(h2, h1, sizeof(h1));
    if (std::memcmp(h2, data.data() + payloadLen, 4) != 0) return false;
    payloadOut.assign(data.begin(), data.begin() + payloadLen);
    return true;
}

bool isValidBase58Check(const std::string& s)
{
    std::vector<std::uint8_t> payload;
    return decodeBase58Check(s, payload);
}

bool isValidBech32(const std::string& s)
{
    if (s.size() < 8 || s.size() > 200) return false;

    // Bech32 forbids mixed case; normalize to lower for verification after that check.
    bool hasLower = false, hasUpper = false;
    for (char c : s) {
        if (c >= 'a' && c <= 'z') hasLower = true;
        else if (c >= 'A' && c <= 'Z') hasUpper = true;
        if (c < 33 || c > 126) return false;  // printable ASCII only
    }
    if (hasLower && hasUpper) return false;

    std::string lower(s);
    std::transform(lower.begin(), lower.end(), lower.begin(),
                   [](unsigned char c) { return static_cast<char>(std::tolower(c)); });

    const std::size_t sep = lower.rfind('1');
    if (sep == std::string::npos || sep == 0) return false;       // need a non-empty HRP
    const std::string hrp = lower.substr(0, sep);
    const std::string data = lower.substr(sep + 1);
    if (data.size() < 6) return false;                            // 6-char checksum minimum

    std::vector<int> values;
    values.reserve(data.size());
    for (char c : data) {
        const char* p = std::strchr(kBech32, c);
        if (p == nullptr) return false;
        values.push_back(static_cast<int>(p - kBech32));
    }

    std::vector<int> combined = bech32HrpExpand(hrp);
    combined.insert(combined.end(), values.begin(), values.end());
    return bech32Polymod(combined) == 1;  // original Bech32 constant (Sapling, not Bech32m)
}

std::string bech32Hrp(const std::string& s)
{
    if (!isValidBech32(s)) return {};
    // isValidBech32 already rejected mixed case and guaranteed a non-empty HRP before the
    // final '1' separator, so lower-casing and splitting there recovers the HRP verbatim.
    std::string lower(s);
    std::transform(lower.begin(), lower.end(), lower.begin(),
                   [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
    const std::size_t sep = lower.rfind('1');
    if (sep == std::string::npos) return {};
    return lower.substr(0, sep);
}

bool isTransparentAddress(const std::string& s)
{
    std::vector<std::uint8_t> payload;
    // version byte (1) + hash160 (20) = 21 bytes, checksum stripped. Covers P2PKH ('R', v60) and
    // P2SH/multisig ('b', v85); the daemon vets the exact version byte for the active network.
    return decodeBase58Check(s, payload) && payload.size() == 21;
}

bool isShieldedAddress(const std::string& s)
{
    const std::string hrp = bech32Hrp(s);
    return hrp == "zs"                 // mainnet Sapling payment address
        || hrp == "ztestsapling"       // testnet
        || hrp == "zregtestsapling";   // regtest
}

bool isValidRecipientAddress(const std::string& s)
{
    return isTransparentAddress(s) || isShieldedAddress(s);
}

}  // namespace util
}  // namespace dragonx
