/*******************************************************************************
 * Copyright (c) 2025, Jan Koester jan.koester@gmx.net
 * All rights reserved.
 * * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 * Redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer.
 * Redistributions in binary form must reproduce the above copyright
 * notice, this list of conditions and the following disclaimer in the
 * documentation and/or other materials provided with the distribution.
 * Neither the name of the <organization> nor the
 * names of its contributors may be used to endorse or promote products
 * derived from this software without specific prior written permission.
 * * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> BE LIABLE FOR ANY
 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *******************************************************************************/

#include <iostream>
#include <iomanip>
#include <algorithm>
#include <stdexcept>
#include <sstream>
#include <fstream>
#include <cstring>

#include "random.h"
#include "rsa.h"

#if defined(_MSC_VER) && !defined(__clang__)
#define  __restrict__ __restrict
#endif // !__restrict__


namespace netplus {
    // Forward declaration for parseDERAndLoadRSA
    static bool parseDERAndLoadRSA(const std::vector<uint8_t>& derData, netplus::rsa& out_rsa);

    netplus::rsa::bigInt::bigInt(const bigInt& other)
    : used(other.used),
      capacity(other.capacity)
    {
        // 1. Allocate a new buffer of the same capacity
        if (capacity > 0) {
            data = std::make_unique<uint32_t[]>(capacity);

            // 2. Perform a fast memory copy of the words
            // We only strictly need to copy up to 'other.used',
            // but copying the full capacity ensures 'ghost data' is consistent.
            std::copy(other.data.get(), other.data.get() + capacity, data.get());
        } else {
            data = nullptr;
        }
    }

    // Move constructor for bigInt
    netplus::rsa::bigInt::bigInt(bigInt&& other) noexcept
    : data(std::move(other.data)),
      used(other.used),
      capacity(other.capacity)
    {
        other.used = 0;
        other.capacity = 0;
    }

        // Assignment from buffer
    netplus::rsa::rsa(const std::vector<uint8_t>& buffer) {
        if (buffer.empty()) {
            throw std::runtime_error("Buffer is empty");
        }

        if (!parseDERAndLoadRSA(buffer, *this)) {
            throw std::runtime_error("Failed to parse RSA key from DER buffer");
        }
    }

    // Assignment from filepath
    netplus::rsa::rsa(const std::string& filepath) {
        std::ifstream file(filepath, std::ios::binary | std::ios::ate);
        if (!file.is_open()) {
            throw std::runtime_error("Cannot open file: " + filepath);
        }

        std::streamsize fileSize = file.tellg();
        if (fileSize <= 0) {
            throw std::runtime_error("File is empty: " + filepath);
        }

        file.seekg(0, std::ios::beg);
        std::vector<uint8_t> derData(static_cast<size_t>(fileSize));
        
        if (!file.read(reinterpret_cast<char*>(derData.data()), fileSize)) {
            throw std::runtime_error("Failed to read file: " + filepath);
        }

        if (!parseDERAndLoadRSA(derData, *this)) {
            throw std::runtime_error("Failed to parse RSA key from file: " + filepath);
        }
    }

    netplus::rsa::bigInt& netplus::rsa::bigInt::operator=(const bigInt& other) {
        if (this == &other) return *this; // Self-assignment check

        // 1. Reallocate if the current capacity is too small
        if (capacity < other.capacity) {
            data = std::make_unique<uint32_t[]>(other.capacity);
            capacity = other.capacity;
        }

        // 2. Copy the state
        used = other.used;
        std::copy(other.data.get(), other.data.get() + used, data.get());

        // 3. Zero out the rest of the capacity for safety
        if (capacity > used) {
            std::fill(data.get() + used, data.get() + capacity, 0);
        }

        return *this;
    }

    // Move assignment for bigInt
    netplus::rsa::bigInt& netplus::rsa::bigInt::operator=(bigInt&& other) noexcept {
        if (this == &other) return *this;
        
        data = std::move(other.data);
        used = other.used;
        capacity = other.capacity;
        
        other.used = 0;
        other.capacity = 0;
        
        return *this;
    }

    void rsa::bigInt::shiftLeft(size_t n) {
        if (n == 0 || isZero()) return;

        size_t wordShift = n / 32;
        size_t bitShift  = n % 32;

        size_t newUsed = used + wordShift + (bitShift ? 1 : 0);
        reserve(newUsed);

        // temp copy of current limbs
        std::vector<uint32_t> old(used);
        std::copy(data.get(), data.get() + used, old.data());

        // zero destination
        std::fill(data.get(), data.get() + newUsed, 0);

        uint32_t carry = 0;
        for (size_t i = 0; i < old.size(); ++i) {
            uint64_t v = old[i];
            if (bitShift) {
                uint64_t shifted = (v << bitShift) | carry;
                data[i + wordShift] = (uint32_t)shifted;
                carry = (uint32_t)(v >> (32 - bitShift));
            } else {
                data[i + wordShift] = (uint32_t)v;
            }
        }
        if (bitShift) data[old.size() + wordShift] = carry;

        used = newUsed;
        while (used > 1 && data[used - 1] == 0) used--;
    }

    bool rsa::bigInt::isBitSet(size_t n) const {
        size_t wordIdx = n / 32;
        if (wordIdx >= used) return false;
        return (data[wordIdx] & (1U << (n % 32))) != 0;
    }

    void netplus::rsa::bigInt::reserve(size_t new_cap) {
        if (new_cap <= capacity) return;
        auto new_data = std::make_unique<uint32_t[]>(new_cap);
        if (data) {
            std::copy(data.get(), data.get() + used, new_data.get());
        }
        data = std::move(new_data);
        capacity = new_cap;
    }

    void netplus::rsa::bigInt::shiftRight(size_t n) {
        if (n == 0 || isZero()) return;
        size_t wordShift = n / 32;
        size_t bitShift = n % 32;

        if (wordShift >= used) {
            used = 1; data[0] = 0;
            return;
        }

        for (size_t i = 0; i < used - wordShift; ++i) {
            uint32_t val = data[i + wordShift];
            if (bitShift > 0 && i + wordShift + 1 < used) {
                val = (val >> bitShift) | (data[i + wordShift + 1] << (32 - bitShift));
            } else {
                val >>= bitShift;
            }
            data[i] = val;
        }
        used -= wordShift;
        while (used > 1 && data[used - 1] == 0) used--;
    }

    void rsa::bigInt::divideByTwo() {
        if (isZero()) return;
        for (size_t i = 0; i < used; ++i) {
            data[i] >>= 1;
            if (i + 1 < used && (data[i + 1] & 1)) {
                data[i] |= 0x80000000;
            }
        }
        if (used > 1 && data[used - 1] == 0) used--;
    }

    // Sets the bigInt value from a standard 64-bit integer
    void netplus::rsa::bigInt::fromUint64(uint64_t val) {
        reserve(2);
        data[0] = (uint32_t)(val & 0xFFFFFFFFu);
        data[1] = (uint32_t)(val >> 32);
        used = (data[1] ? 2 : 1);   // zero becomes used=1, data[0]=0
    }

    // Simple equality check for assertions without needing strings
    bool netplus::rsa::bigInt::equalsUint64(uint64_t val) const {
        if (val == 0) return isZero();
        if (used > 2) return false;
        uint64_t current = data[0];
        if (used == 2) current |= (static_cast<uint64_t>(data[1]) << 32);
        return current == val;
    }

    std::string netplus::rsa::bigInt::toHexString() const {
        if (used == 0 || (used == 1 && data[0] == 0)) return "0";

        std::stringstream ss;
        ss << std::hex << std::uppercase;

        // First word no leading zeros
        ss << data[used - 1];

        // Middle words padded to 8 chars
        for (int i = (int)used - 2; i >= 0; --i) {
            ss << std::setfill('0') << std::setw(8) << data[i];
        }

        return ss.str();
    }

    rsa::bigInt rsa::bigInt::operator*(const rsa::bigInt& other) const {
        bigInt res(this->capacity + other.capacity);
        rsa::multiply(*this, other, res);
        return res;
    }

    rsa::bigInt rsa::bigInt::operator%(const rsa::bigInt& mod) const {
        bigInt q(this->capacity), r(mod.capacity);
        rsa::divide(*this, mod, q, r);
        return r;
    }


    rsa::rsa() : n(64), e(64), d(64) {
        // Explicitly initialize bigInt members with capacity 64
    }

    // Copy constructor
    rsa::rsa(const rsa& src)
    : n(src.n), e(src.e), d(src.d)
    {
    }

    // Move constructor
    rsa::rsa(rsa&& src) noexcept
    : n(std::move(src.n)), e(std::move(src.e)), d(std::move(src.d))
    {
    }

    // Copy assignment
    rsa& rsa::operator=(const rsa& src) {
        if (this != &src) {
            n = src.n;
            e = src.e;
            d = src.d;
        }
        return *this;
    }

    // Move assignment
    rsa& rsa::operator=(rsa&& src) noexcept {
        if (this != &src) {
            n = std::move(src.n);
            e = std::move(src.e);
            d = std::move(src.d);
        }
        return *this;
    }

    // Helper function to parse a DER length field
    static bool parseDERLength(const uint8_t* data, size_t maxLen, size_t& outLength, size_t& outBytesRead) {
        if (maxLen < 1) return false;
        
        outLength = data[0];
        outBytesRead = 1;
        
        // Long form
        if (outLength & 0x80) {
            int numOctets = static_cast<int>(outLength & 0x7F);
            if (numOctets <= 0 || numOctets > 4) return false;
            if (maxLen < 1 + static_cast<size_t>(numOctets)) return false;
            
            outLength = 0;
            for (int i = 0; i < numOctets; ++i) {
                outLength = (outLength << 8) | data[1 + i];
            }
            outBytesRead = 1 + numOctets;
        }
        
        return true;
    }

    // Helper function to parse a DER INTEGER
    static bool parseDERInteger(const uint8_t* data, size_t maxLen, std::vector<uint8_t>& outBytes, size_t& outBytesRead) {
        if (maxLen < 2 || data[0] != 0x02) {  // 0x02 = INTEGER tag
            return false;
        }
        
        size_t length = 0;
        size_t lenBytesRead = 0;
        if (!parseDERLength(data + 1, maxLen - 1, length, lenBytesRead)) {
            return false;
        }
        
        if (1 + lenBytesRead + length > maxLen) {
            return false;
        }
        
        const uint8_t* intData = data + 1 + lenBytesRead;
        outBytes.clear();
        
        // Skip leading zero byte if present (DER encoding convention for non-negative integers)
        if (length > 0 && intData[0] == 0x00) {
            outBytes.assign(intData + 1, intData + length);
        } else {
            outBytes.assign(intData, intData + length);
        }
        
        outBytesRead = 1 + lenBytesRead + length;
        return true;
    }

    // Helper function to check if buffer contains RSA OID (1.2.840.113549.1.1.1)
    static bool isRsaOid(const uint8_t* data, size_t len) {
        // RSA OID: 06 09 2a 86 48 86 f7 0d 01 01 01
        static const uint8_t rsaOid[] = {0x06, 0x09, 0x2a, 0x86, 0x48, 0x86, 0xf7, 0x0d, 0x01, 0x01, 0x01};
        if (len < sizeof(rsaOid)) return false;
        return std::memcmp(data, rsaOid, sizeof(rsaOid)) == 0;
    }

    // Helper function to extract PKCS#1 key from PKCS#8 envelope
    static bool extractPKCS1FromPKCS8(const std::vector<uint8_t>& pkcs8Data, std::vector<uint8_t>& pkcs1Data) {
        // PKCS#8 PrivateKeyInfo structure:
        // SEQUENCE {
        //   INTEGER version (0)
        //   SEQUENCE { OID algorithm, NULL parameters }
        //   OCTET STRING privateKey  <-- this contains the PKCS#1 key
        // }
        
        if (pkcs8Data.size() < 20) return false;
        
        const uint8_t* data = pkcs8Data.data();
        size_t size = pkcs8Data.size();
        size_t cursor = 0;
        
        // Outer SEQUENCE
        if (data[cursor] != 0x30) return false;
        cursor++;
        
        size_t outerLen = 0, lenBytes = 0;
        if (!parseDERLength(data + cursor, size - cursor, outerLen, lenBytes)) return false;
        cursor += lenBytes;
        
        // Version INTEGER (should be 0)
        if (data[cursor] != 0x02) return false;
        std::vector<uint8_t> version;
        size_t intBytes = 0;
        if (!parseDERInteger(data + cursor, size - cursor, version, intBytes)) return false;
        cursor += intBytes;
        
        // Algorithm SEQUENCE
        if (data[cursor] != 0x30) return false;
        cursor++;
        
        size_t algLen = 0;
        if (!parseDERLength(data + cursor, size - cursor, algLen, lenBytes)) return false;
        cursor += lenBytes;
        
        // Check for RSA OID within algorithm sequence
        if (!isRsaOid(data + cursor, algLen)) return false;
        cursor += algLen;  // skip algorithm sequence content
        
        // Private key OCTET STRING
        if (data[cursor] != 0x04) return false;
        cursor++;
        
        size_t keyLen = 0;
        if (!parseDERLength(data + cursor, size - cursor, keyLen, lenBytes)) return false;
        cursor += lenBytes;
        
        if (cursor + keyLen > size) return false;
        
        // Extract the PKCS#1 key
        pkcs1Data.assign(data + cursor, data + cursor + keyLen);
        return true;
    }

    // Helper function to parse DER data and load RSA key
    static bool parseDERAndLoadRSA(const std::vector<uint8_t>& derData, netplus::rsa& out_rsa) {
        if (derData.empty()) {
            return false;
        }
        
        // First, try to unwrap PKCS#8 envelope if present
        std::vector<uint8_t> pkcs1Data;
        if (extractPKCS1FromPKCS8(derData, pkcs1Data)) {
            // Recursively parse the extracted PKCS#1 key
            return parseDERAndLoadRSA(pkcs1Data, out_rsa);
        }
        
        const uint8_t* data = derData.data();
        size_t size = derData.size();
        size_t cursor = 0;
        
        // Must start with SEQUENCE (0x30)
        if (size < 2 || data[cursor] != 0x30) {
            return false;
        }
        cursor++;
        
        // Parse SEQUENCE length
        size_t seqLength = 0;
        size_t lenBytesRead = 0;
        if (!parseDERLength(data + cursor, size - cursor, seqLength, lenBytesRead)) {
            return false;
        }
        cursor += lenBytesRead;
        
        if (cursor + seqLength > size) {
            return false;
        }
        
        std::vector<uint8_t> n, e, d;
        
        // Try to parse as PKCS#1 RSAPrivateKey first
        size_t tempCursor = cursor;
        bool isPKCS1Private = false;
        
        // Check for version (INTEGER 0)
        if (tempCursor < cursor + seqLength && data[tempCursor] == 0x02) {
            std::vector<uint8_t> version;
            size_t intBytesRead = 0;
            if (parseDERInteger(data + tempCursor, seqLength - (tempCursor - cursor), version, intBytesRead)) {
                tempCursor += intBytesRead;
                
                // Try to parse n, e, d
                if (parseDERInteger(data + tempCursor, seqLength - (tempCursor - cursor), n, intBytesRead)) {
                    tempCursor += intBytesRead;
                    
                    if (parseDERInteger(data + tempCursor, seqLength - (tempCursor - cursor), e, intBytesRead)) {
                        tempCursor += intBytesRead;
                        
                        if (parseDERInteger(data + tempCursor, seqLength - (tempCursor - cursor), d, intBytesRead)) {
                            isPKCS1Private = true;
                        }
                    }
                }
            }
        }
        
        if (isPKCS1Private && !n.empty() && !e.empty() && !d.empty()) {
            // Successfully parsed PKCS#1 private key
            out_rsa.setRsaKeyFromRaw(n, e, d);
            return true;
        }
        
        // Try to parse as public key (might be SubjectPublicKeyInfo with BIT STRING, or just RSAPublicKey SEQUENCE)
        if (data[cursor] == 0x30) {
            // Could be either RSAPublicKey or algorithm identifier
            cursor++;
            size_t innerSeqLength = 0;
            if (!parseDERLength(data + cursor, size - cursor, innerSeqLength, lenBytesRead)) {
                return false;
            }
            cursor += lenBytesRead;
            
            if (cursor + innerSeqLength > size) {
                return false;
            }
            
            // Try to parse modulus and exponent
            size_t pubCursor = cursor;
            size_t intBytesRead = 0;
            if (parseDERInteger(data + pubCursor, innerSeqLength - (pubCursor - cursor), n, intBytesRead)) {
                pubCursor += intBytesRead;
                
                if (parseDERInteger(data + pubCursor, innerSeqLength - (pubCursor - cursor), e, intBytesRead)) {
                    if (!n.empty() && !e.empty()) {
                        // Create a public-only key
                        d.clear();  // No private exponent
                        out_rsa.setRsaKeyFromRaw(n, e, d);
                        return true;
                    }
                }
            }
        }
        
        return false;
    }

    // R = 2^(32*nwords), so R2 = R*R mod N
    rsa::bigInt rsa::calculateR2Mod(const bigInt& mod) {
        bigInt r(mod.used * 4 + 2);
        r.fromUint64(1);
        r.shiftLeft(mod.used * 64); // 2*(32*nwords)
        bigInt q(r.used + 1), rem(mod.used + 1);
        divide(r, mod, q, rem);
        return rem;
    }

    void rsa::add(const bigInt& a, const bigInt& b, bigInt& res) {
        uint64_t carry = 0;
        const size_t au = a.used;
        const size_t bu = b.used;
        size_t max_len = std::max(au, bu);

        res.reserve(max_len + 1);
        res.used = 0;

        const uint32_t* ap = a.data.get();
        const uint32_t* bp = b.data.get();
        uint32_t* rp = res.data.get();

        for (size_t i = 0; i < max_len || carry; ++i) {
            uint64_t sum = carry
                + (i < au ? ap[i] : 0)
                + (i < bu ? bp[i] : 0);

            rp[i] = (uint32_t)sum;
            carry = sum >> 32;
            res.used++;
        }

        if (res.used == 0) { res.used = 1; rp[0] = 0; }
    }


    void rsa::multiply(const bigInt& a, const bigInt& b, bigInt& res) {
        const uint32_t* __restrict__ ap = a.data.get();
        const uint32_t* __restrict__ bp = b.data.get();
        uint32_t* __restrict__ rp = res.data.get();
        const size_t au = a.used;
        const size_t bu = b.used;

        // Zero only the output region we'll actually write
        std::memset(rp, 0, (au + bu) * sizeof(uint32_t));

        for (size_t i = 0; i < au; ++i) {
            uint64_t carry = 0;
            const uint64_t ai = (uint64_t)ap[i];
            for (size_t j = 0; j < bu; ++j) {
                uint64_t cur = (uint64_t)rp[i + j] + ai * (uint64_t)bp[j] + carry;
                rp[i + j] = (uint32_t)cur;
                carry = cur >> 32;
            }

            size_t k = i + bu;
            uint64_t cur = (uint64_t)rp[k] + carry;
            rp[k] = (uint32_t)cur;
            carry = cur >> 32;

            while (carry) {
                ++k;
                uint64_t cur2 = (uint64_t)rp[k] + carry;
                rp[k] = (uint32_t)cur2;
                carry = cur2 >> 32;
            }
        }

        res.used = au + bu;
        while (res.used > 1 && rp[res.used - 1] == 0) res.used--;
    }


    int rsa::compare(const bigInt& a, const bigInt& b) {
        if (a.used != b.used)
            return (a.used > b.used) ? 1 : -1;

        const uint32_t* ap = a.data.get();
        const uint32_t* bp = b.data.get();
        for (int i = (int)a.used - 1; i >= 0; --i) {
            if (ap[i] != bp[i])
                return (ap[i] > bp[i]) ? 1 : -1;
        }
        return 0;
    }

    size_t rsa::bigInt::bitLength() const {
        if (isZero()) return 0;
        size_t bits = (used - 1) * 32;
        uint32_t lastWord = data[used - 1];
        while (lastWord > 0) {
            lastWord >>= 1;
            bits++;
        }
        return bits;
    }

    void rsa::subtract(const bigInt& a, const bigInt& b, bigInt& res) {
        int64_t borrow = 0;
        const size_t au = a.used;
        const size_t bu = b.used;
        res.reserve(au);
        const uint32_t* ap = a.data.get();
        const uint32_t* bp = b.data.get();
        uint32_t* rp = res.data.get();
        for (size_t i = 0; i < au; ++i) {
            int64_t sub = (int64_t)ap[i] - (i < bu ? bp[i] : 0) - borrow;
            if (sub < 0) {
                sub += 0x100000000ULL;
                borrow = 1;
            } else {
                borrow = 0;
            }
            rp[i] = (uint32_t)sub;
        }
        res.used = au;
        while (res.used > 1 && rp[res.used - 1] == 0) res.used--;
    }

    netplus::rsa::bigInt netplus::rsa::modPow(const netplus::rsa::bigInt& base,
                                    const netplus::rsa::bigInt& exp,
                                    const netplus::rsa::bigInt& mod)
    {
        using bigInt = netplus::rsa::bigInt;

        // Handle mod == 0 or 1 edge cases
        if (mod.used == 1 && mod.data[0] == 0) {
            throw std::runtime_error("modPow: modulus is zero");
        }
        if (mod.used == 1 && mod.data[0] == 1) {
            bigInt z(1U, 1);
            z.data[0] = 0;
            z.used = 1;
            return z;
        }

        // exp == 0 => 1 mod mod
        if (exp.isZero()) {
            bigInt one(1U, 1);
            bigInt q(one.capacity), r(mod.capacity);
            netplus::rsa::divide(one, mod, q, r);
            return r;
        }

        // Montgomery requires odd modulus
        if (mod.isEven()) {
            // Fallback to classic square-and-multiply if needed (not expected for RSA)
            bigInt result(1U, mod.capacity);
            bigInt a = base % mod;

            for (size_t i = exp.bitLength(); i-- > 0; ) {
                result = (result * result) % mod;
                if (exp.isBitSet(i)) result = (result * a) % mod;
            }
            return result;
        }

        // n_prime = -N^{-1} mod 2^32 (uses least significant limb)
        uint32_t n_prime = netplus::rsa::calculateNPrime(mod.data[0]);

        // Helper: compute (2^(64*mod.used)) mod mod == R^2 mod N
        auto calculateR2Mod = [&](const bigInt& N) -> bigInt {
            bigInt r(N.used * 4 + 2);
            r.fromUint64(1);
            r.shiftLeft(N.used * 64);

            bigInt q(r.used + 1), rem(N.used + 2);
            netplus::rsa::divide(r, N, q, rem);
            return rem;
        };

        size_t n = mod.used;

        // Pre-allocate ALL scratch buffers once — no allocations inside the loop
        bigInt scratch(2 * n + 2);   // for montgomeryMultiply_into product
        bigInt temp(n + 2);          // for montgomeryMultiply_into output

        // Convert into Montgomery domain
        bigInt a = base % mod;
        bigInt R2 = calculateR2Mod(mod);

        bigInt a_bar(n + 2);
        montgomeryMultiply_into(a, R2, mod, n_prime, scratch, a_bar);

        // res_bar = 1 * R mod N  (Montgomery form of 1)
        bigInt one(1U, 1);
        bigInt res_bar(n + 2);
        montgomeryMultiply_into(one, R2, mod, n_prime, scratch, res_bar);

        // Exponentiation (left-to-right square-and-multiply) — zero allocations
        for (size_t i = exp.bitLength(); i-- > 0; ) {
            // res_bar = montMul(res_bar, res_bar) — square
            montgomeryMultiply_into(res_bar, res_bar, mod, n_prime, scratch, temp);
            std::swap(res_bar, temp);

            if (exp.isBitSet(i)) {
                // res_bar = montMul(res_bar, a_bar) — multiply
                montgomeryMultiply_into(res_bar, a_bar, mod, n_prime, scratch, temp);
                std::swap(res_bar, temp);
            }
        }

        // Convert out of Montgomery domain: montMul(res_bar, 1) = result
        montgomeryMultiply_into(res_bar, one, mod, n_prime, scratch, temp);
        return temp;
    }

    void rsa::divide(const bigInt& dividend, const bigInt& divisor, bigInt& quotient, bigInt& remainder) {
        remainder = dividend;
        quotient.reserve(dividend.used);
        std::fill(quotient.data.get(), quotient.data.get() + quotient.capacity, 0);
        quotient.used = 1;

        if (compare(dividend, divisor) < 0) return;

        bigInt tempDiv = divisor;
        int shift = (int)dividend.bitLength() - (int)divisor.bitLength();
        tempDiv.shiftLeft(shift);

        for (int i = 0; i <= shift; ++i) {
            if (compare(remainder, tempDiv) >= 0) {
                subtract(remainder, tempDiv, remainder);
                // Set the bit in the quotient
                int qBit = shift - i;
                quotient.data[qBit / 32] |= (1U << (qBit % 32));
                if ((size_t)(qBit / 32) >= quotient.used) quotient.used = (qBit / 32) + 1;
            }
            tempDiv.shiftRight(1);
        }
        while (quotient.used > 1 && quotient.data[quotient.used - 1] == 0) quotient.used--;
    }

    std::vector<uint8_t> bigIntToBytes(const netplus::rsa::bigInt& n, size_t targetLen) {
        std::vector<uint8_t> result(targetLen, 0);
        // Copy the internal 32-bit words into 8-bit bytes (Big Endian)
        for (size_t i = 0; i < n.used; ++i) {
            uint32_t word = n.data[i];
            for (int j = 0; j < 4; ++j) {
                size_t bytePos = targetLen - 1 - (i * 4 + j);
                if (bytePos < targetLen) {
                    result[bytePos] = (word >> (j * 8)) & 0xFF;
                }
            }
        }
        return result;
    }

    void netplus::rsa::generateSecureRandom(bigInt& n, size_t bits) {

        size_t words = (bits + 31) / 32;

        n.reserve(words);
        n.used = words;

        fillRandomBytes(reinterpret_cast<uint8_t*>(n.data.get()),
                        words * sizeof(uint32_t));

        n.data[words - 1] |= (1U << ((bits - 1) % 32));

        uint32_t excess = (uint32_t)(words * 32 - bits);
        if (excess > 0) {
            uint32_t mask = 0xFFFFFFFFu >> excess;
            n.data[words - 1] &= mask;
            n.data[words - 1] |= (1U << ((bits - 1) % 32));
        }
        n.data[0] |= 1U;
    }


    rsa::bigInt rsa::gcd(bigInt a, bigInt b) {
        while (!b.isZero()) {
            bigInt q(a.capacity), r(a.capacity);
            divide(a, b, q, r);
            a = b;
            b = r;
        }
        return a;
    }

    rsa::bigInt rsa::modInverse(const bigInt& k, const bigInt& pm) {
        // 1. Initialize local constants properly to avoid ambiguity
        bigInt one((uint32_t)1, (size_t)1);

        bigInt u = k;
        bigInt v = pm;

        // Ensure capacity is sufficient for intermediate additions
        size_t cap = pm.capacity + 1;
        bigInt x1((size_t)cap);
        bigInt x2((size_t)cap);

        x1.data[0] = 1;
        x1.used = 1;
        // x2 is initialized to 0 by the constructor

        // 2. Binary Extended Euclidean Algorithm
        while (compare(u, one) != 0 && compare(v, one) != 0) {
            while (u.isEven()) {
                u.divideByTwo();
                if (x1.isEven()) {
                    x1.divideByTwo();
                } else {
                    add(x1, pm, x1);
                    x1.divideByTwo();
                }
            }
            while (v.isEven()) {
                v.divideByTwo();
                if (x2.isEven()) {
                    x2.divideByTwo();
                } else {
                    add(x2, pm, x2);
                    x2.divideByTwo();
                }
            }
            if (compare(u, v) >= 0) {
                subtract(u, v, u);
                if (compare(x1, x2) < 0) add(x1, pm, x1);
                subtract(x1, x2, x1);
            } else {
                subtract(v, u, v);
                if (compare(x2, x1) < 0) add(x2, pm, x2);
                subtract(x2, x1, x2);
            }
        }

        return (compare(u, one) == 0) ? x1 : x2;
    }

    bool rsa::isProbablyPrime(const bigInt& n, int k) {
        bigInt one(1U, n.capacity);
        bigInt two(2U, n.capacity);
        bigInt three(3U, n.capacity);

        // nm1 = n - 1
        bigInt nm1(n.capacity);
        subtract(n, one, nm1);

        // trivial checks
        if (compare(n, one) <= 0) return false;
        if (n.isEven()) return compare(n, two) == 0;

        // ------------------------------------------------------------
        // Lambda: randomBigIntBits(out,bits)
        // ------------------------------------------------------------
        auto randomBigIntBits = [&](bigInt& out, size_t bits) {
            size_t words = (bits + 31) / 32;
            out.reserve(words);
            out.used = words;

            // fill random words
            for (size_t i = 0; i < words; ++i) {
                fillRandomBytes(reinterpret_cast<uint8_t*>(&out.data[i]),
                                sizeof(uint32_t));
            }

            // mask unused high bits
            size_t extra = words * 32 - bits;
            if (extra > 0) {
                uint32_t mask = 0xFFFFFFFFu >> extra;
                out.data[words - 1] &= mask;
            }

            // optional normalize if you have it:
            // out.normalize();
        };

        // ------------------------------------------------------------
        // Lambda: randomBelow(limit)   -> uniform r in [0..limit-1]
        // ------------------------------------------------------------
        auto randomBelow = [&](const bigInt& limit) -> bigInt {
            size_t bits = limit.bitLength();
            bigInt r(limit.capacity);

            for (;;) {
                randomBigIntBits(r, bits);

                if (compare(r, limit) < 0)
                    return r;
            }
        };

        // ------------------------------------------------------------
        // Lambda: randomWitnessA(n) -> uniform a in [2..n-2]
        // ------------------------------------------------------------
        auto randomWitnessA = [&]() -> bigInt {

            // nMinusThree = n - 3
            bigInt nMinusThree(n.capacity);
            subtract(n, three, nMinusThree);   // n-3

            // r ∈ [0..n-4]
            bigInt r = randomBelow(nMinusThree);

            // a = r + 2  => [2..n-2]
            bigInt a(n.capacity);
            add(r, two, a);
            return a;
        };

        // ------------------------------------------------------------
        // factor n-1 = 2^s * d
        // ------------------------------------------------------------
        bigInt d = nm1;
        size_t s = 0;
        while (d.isEven()) {
            d.divideByTwo();
            s++;
        }

        // ------------------------------------------------------------
        // Miller-Rabin rounds
        // ------------------------------------------------------------
        for (int i = 0; i < k; i++) {

            bigInt a = randomWitnessA();

            bigInt x = modPow(a, d, n);

            // x == 1 or x == n-1
            if ((x.used == 1 && x.data[0] == 1) || compare(x, nm1) == 0)
                continue;

            bool composite = true;

            for (size_t r = 1; r < s; r++) {
                x = modPow(x, two, n);

                if (compare(x, nm1) == 0) {
                    composite = false;
                    break;
                }
            }

            if (composite) return false;
        }

        return true;
    }

    void rsa::findPrime(bigInt& p, size_t bits) {
        // Correctly initialize 'two' to avoid ambiguous constructor call
        bigInt two(2U, 1);

        // Initial random start (ensured odd by generateSecureRandom)
        generateSecureRandom(p, bits);

        // Miller-Rabin Primality Test Loop
        while (!isProbablyPrime(p, 10)) {
            // We can add directly to p to avoid creating 'next_p' every time
            // if your add function supports add(a, b, a)
            bigInt next_p(p.capacity);
            add(p, two, next_p);
            p = next_p;
        }
    }

    void rsa::generateKeys(size_t bit_count) {
        // bit_count for p and q should be half of the target N bit length
        bigInt p(bit_count), q(bit_count);

        findPrime(p, bit_count);
        findPrime(q, bit_count);

        // n = p * q
        n.reserve(p.used + q.used);
        multiply(p, q, n);

        // Constants for subtraction
        bigInt one(1U, 1);
        bigInt two(2U, 1);

        // phi = (p-1)(q-1)
        bigInt pm1(p.capacity), qm1(q.capacity), phi(n.capacity);
        subtract(p, one, pm1);
        subtract(q, one, qm1);
        multiply(pm1, qm1, phi);

        // Standard public exponent e = 65537
        // Using 65537U is more secure and faster than 17
        e = bigInt(65537U, 8);

        // Ensure e is coprime to phi
        while (true) {
            bigInt common = gcd(e, phi);
            // Check if gcd == 1
            if (common.used == 1 && common.data[0] == 1) {
                break;
            }
            bigInt next_e(e.capacity);
            add(e, two, next_e);
            e = next_e;
        }

        // Calculate d = e^-1 mod phi
        this->d = modInverse(e, phi);

        // Optional: Verification logic without toString()
        /*
        bigInt check(n.capacity * 2);
        multiply(e, this->d, check);
        bigInt q_c(check.capacity), r_c(phi.capacity);
        divide(check, phi, q_c, r_c);
        if (r_c.used == 1 && r_c.data[0] == 1) {
            // Verification Success: (e*d % phi) == 1
        }
        */
    }

    void netplus::rsa::setRsaKeyFromRaw(
            const std::vector<uint8_t>& nBE,
            const std::vector<uint8_t>& eBE,
            const std::vector<uint8_t>& dBE
        ) {

            auto bigIntFromBytesBE = [&](const uint8_t* bytes, size_t len) -> rsa::bigInt {
                // Determine how many 32-bit words we need
                size_t words = (len + 3) / 4;
                rsa::bigInt x(words + 2);
                x.used = words;

                // Fill words from the end of the byte array
                for (size_t i = 0; i < len; ++i) {
                    size_t bytePos = len - 1 - i;
                    size_t wordIdx = i / 4;
                    size_t shift = (i % 4) * 8;
                    x.data[wordIdx] |= (static_cast<uint32_t>(bytes[bytePos]) << shift);
                }
                return x;
            };

            n = bigIntFromBytesBE(nBE.data(), nBE.size());
            e = bigIntFromBytesBE(eBE.data(), eBE.size());
            d = bigIntFromBytesBE(dBE.data(), dBE.size());
        }

    rsa::bigInt rsa::encrypt(const bigInt& m) { return modPow(m, e, n); }
    rsa::bigInt rsa::decrypt(const bigInt& c) { return modPow(c, d, n); }

    // This is used to "clear" the lower bits during reduction
    uint32_t rsa::calculateNPrime(uint32_t n0) {
        uint32_t inv = 1;
        for (int i = 0; i < 32; i++) {
            inv *= 2 - n0 * inv;
        }
        return -inv;
    }

    // 2. Calculate R mod N (where R = 2^(32 * mod.used))
    rsa::bigInt rsa::calculateRMod(const bigInt& mod) {
        bigInt r(mod.used * 2 + 1);
        r.fromUint64(1);

        // Shift 1 left by (32 * mod.used) bits, then mod N
        // This is equivalent to R mod N
        r.shiftLeft(mod.used * 32);

        bigInt q(r.used), remainder(mod.used);
        divide(r, mod, q, remainder);
        return remainder;
    }

    // 3. Montgomery Multiplication: (a * b * R^-1) mod N
    // This is the core "speed" function that avoids long division.
    rsa::bigInt rsa::montgomeryMultiply(const bigInt& a, const bigInt& b, const bigInt& mod, uint32_t n_prime) {
        size_t n = mod.used;
        bigInt res(2 * n + 1);
        bigInt final_res(n + 1);

        montgomeryMultiply_into(a, b, mod, n_prime, res, final_res);
        return final_res;
    }

    // In-place Montgomery multiply: writes result into |out|, uses |scratch| as temp.
    // Both scratch and out must be pre-allocated with sufficient capacity.
    void rsa::montgomeryMultiply_into(const bigInt& a, const bigInt& b,
                                       const bigInt& mod, uint32_t n_prime,
                                       bigInt& scratch, bigInt& out) {
        const size_t n = mod.used;
        const size_t need = 2 * n + 2;  // extra word for carry propagation

        // Ensure scratch has enough room
        if (scratch.capacity < need) scratch.reserve(need);

        // Perform schoolbook multiplication (scratch = a * b)
        multiply(a, b, scratch);
        // Ensure any limbs beyond the product are zero for carry propagation
        if (scratch.capacity < need) scratch.reserve(need);
        for (size_t i = scratch.used; i < need; ++i)
            scratch.data[i] = 0;

        // Montgomery Reduction — use raw pointers
        uint32_t* __restrict__ sp = scratch.data.get();
        const uint32_t* __restrict__ mp = mod.data.get();

        for (size_t i = 0; i < n; i++) {
            const uint64_t mi = (uint64_t)(sp[i] * n_prime);
            uint64_t carry = 0;
            for (size_t j = 0; j < n; j++) {
                uint64_t cur = (uint64_t)sp[i + j] + mi * (uint64_t)mp[j] + carry;
                sp[i + j] = (uint32_t)cur;
                carry = cur >> 32;
            }
            size_t k = i + n;
            while (carry > 0) {
                uint64_t cur = (uint64_t)sp[k] + carry;
                sp[k] = (uint32_t)cur;
                carry = cur >> 32;
                k++;
            }
        }

        // Shift right by N words (divide by R) — write into out
        if (out.capacity < n + 1) out.reserve(n + 1);
        uint32_t* __restrict__ op = out.data.get();
        std::memcpy(op, sp + n, (n + 1) * sizeof(uint32_t));
        out.used = n + 1;
        while (out.used > 1 && op[out.used - 1] == 0) out.used--;

        // Final step: if out >= mod, subtract mod in-place
        if (compare(out, mod) >= 0) {
            subtract(out, mod, scratch);
            out.used = scratch.used;
            std::memcpy(op, scratch.data.get(), scratch.used * sizeof(uint32_t));
            if (out.capacity > out.used)
                std::memset(op + out.used, 0, (out.capacity - out.used) * sizeof(uint32_t));
        }
    }

    rsa::bigInt rsa::bytesToBigIntBE(const std::vector<uint8_t>& bytes) {
        size_t words = (bytes.size() + 3) / 4;
        netplus::rsa::bigInt out(std::max<size_t>(words + 1, 1));
        out.used = 1;
        out.data[0] = 0;

        for (uint8_t b : bytes) {
            // out = out << 8
            uint32_t carry = 0;
            for (size_t i = 0; i < out.used; ++i) {
                uint32_t newCarry = out.data[i] >> 24;
                out.data[i] = (out.data[i] << 8) | carry;
                carry = newCarry;
            }
            if (carry != 0) {
                out.data[out.used++] = carry;
            }

            // out = out + b
            uint64_t sum = uint64_t(out.data[0]) + b;
            out.data[0] = uint32_t(sum);
            uint32_t c = uint32_t(sum >> 32);

            size_t idx = 1;
            while (c != 0 && idx < out.used) {
                uint64_t s = uint64_t(out.data[idx]) + c;
                out.data[idx] = uint32_t(s);
                c = uint32_t(s >> 32);
                ++idx;
            }
            if (c != 0) {
                out.data[out.used++] = c;
            }
        }

        // trim leading zero limbs (optional safety)
        while (out.used > 1 && out.data[out.used - 1] == 0) out.used--;

        return out;
    }

    rsa::bigInt rsa::bigIntFromBytesBE (const uint8_t* bytes, size_t len){
        size_t words = (len + 3) / 4;
        rsa::bigInt x(words + 1);
        x.used = words;
        std::fill(x.data.get(), x.data.get() + x.capacity, 0);

        for (size_t i = 0; i < len; ++i) {
            uint8_t byte = bytes[len - 1 - i];     // LSB first
            size_t wordIdx = i / 4;
            size_t bitShift = (i % 4) * 8;
            x.data[wordIdx] |= (static_cast<uint32_t>(byte) << bitShift);
        }
        while (x.used > 1 && x.data[x.used - 1] == 0) x.used--;
        return x;
    };

    std::vector<uint8_t> rsa::bigIntToBytesBE(const bigInt& x, size_t outLen){
        std::vector<uint8_t> out(outLen, 0x00);

        bigInt t = x;

        for (size_t i = 0; i < outLen; i++) {
            out[outLen - 1 - i] = uint8_t(t.data[0] & 0xFF);
            t.shiftRight(8);
        }

        return out;
    }

    // Static function to load RSA key from DER-encoded data


};
