/*******************************************************************************
 * 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 <vector>
#include <iostream>
#include <memory>
#include <string>
#include <algorithm>
#include <cstdint>

#pragma once

// Portable limb types: 64-bit on GCC/Clang, 32-bit fallback on MSVC
#if defined(__SIZEOF_INT128__)
    using limb_t  = uint64_t;
    using dlimb_t = __uint128_t;
    #define LIMB_BITS 64
#else
    using limb_t  = uint32_t;
    using dlimb_t = uint64_t;
    #define LIMB_BITS 32
#endif

namespace netplus {
    class rsa {
    public:
        class bigInt {
            public:
                std::unique_ptr<limb_t[]> data;
                size_t used; // number of limb_t words used
                size_t capacity;

                explicit bigInt(size_t cap_words = 64) : used(1), capacity(std::max<size_t>(1, cap_words)) {
                    data = std::make_unique<limb_t[]>(capacity);
                    std::fill(data.get(), data.get() + capacity, limb_t(0));
                }

                bigInt(uint32_t val, size_t cap) : used(1), capacity(std::max<size_t>(1, cap)) {
                    data = std::make_unique<limb_t[]>(capacity);
                    std::fill(data.get(), data.get() + capacity, limb_t(0));
                    data[0] = val;
                }


                bigInt(const bigInt& other);
                bigInt(bigInt&& other) noexcept;

                netplus::rsa::bigInt& operator=(const bigInt& other);
                netplus::rsa::bigInt& operator=(bigInt&& other) noexcept;

                // Helper to check bit state for modPow
                bool isBitSet(size_t bit) const;

                bool isZero() const { return used == 1 && data[0] == 0; }
                bool isEven() const { return (data[0] & 1) == 0; } // safe if used>=1 always


                size_t bitLength() const;

                void shiftLeft(size_t n);
                void reserve(size_t new_cap);
                void shiftRight(size_t n);
                void divideByTwo();
                void fromUint64(uint64_t val);
                bool equalsUint64(uint64_t val) const;
                std::string toHexString() const;

                // Inside struct bigInt
                bigInt operator*(const bigInt& other) const;
                bigInt operator%(const bigInt& mod) const;
            };

        rsa();
        rsa(const rsa& src);
        rsa(rsa&& src) noexcept;
        rsa (const std::string& filepath);
        rsa(const std::vector<uint8_t>& buffer);

        rsa& operator=(const rsa& src);
        rsa& operator=(rsa&& src) noexcept;


        explicit operator bool() const {
            return !(n.isZero() || d.isZero());
        }

        bool operator!() const {
            return !static_cast<bool>(*this);
        }

        // RSA Core Methods
        void generateKeys(size_t digit_count);
        bigInt encrypt(const bigInt& message);
        bigInt decrypt(const bigInt& cipher);

        // Key component accessors (read-only)
        const bigInt& getN() const { return n; }
        const bigInt& getE() const { return e; }
        const bigInt& getD() const { return d; }
        const bigInt& getP() const { return p; }
        const bigInt& getQ() const { return q; }

        // Static Math Engine
        void reserve(size_t new_cap);
        void shiftLeft(size_t n);
        void shiftRight(size_t n);

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

        static void add(const bigInt& a, const bigInt& b, bigInt& res);
        static void subtract(const bigInt& a, const bigInt& b, bigInt& res);
        static void multiply(const bigInt& a, const bigInt& b, bigInt& res);
        static void divide(const bigInt& dividend, const bigInt& divisor, bigInt& quotient, bigInt& remainder);
        static int compare(const bigInt& a, const bigInt& b);
        static bigInt bytesToBigIntBE(const std::vector<uint8_t>& bytes);
        static bigInt bigIntFromBytesBE (const uint8_t* bytes, size_t len);
        static std::vector<uint8_t> bigIntToBytesBE(const bigInt& x, size_t outLen);
        
        static bigInt modInverse(const bigInt& e, const bigInt& phi);

    private:
        bigInt n, e, d, p, q;
        bool isProbablyPrime(const bigInt& n, int k);
        void findPrime(bigInt& p, size_t digits);
        void generateSecureRandom(bigInt& n, size_t requested_digits);
        static bigInt modPow(const bigInt& base, const bigInt& exp, const bigInt& mod);
        static bigInt gcd(bigInt a, bigInt b);
        static limb_t calculateNPrime(limb_t n0);
        static bigInt calculateRMod(const bigInt& mod);
        static bigInt calculateR2Mod(const bigInt& mod);
        static bigInt montgomeryMultiply(const bigInt& a, const bigInt& b, const bigInt& mod, limb_t n_prime);
        static void montgomeryMultiply_into(const bigInt& a, const bigInt& b,
                                             const bigInt& mod, limb_t n_prime,
                                             bigInt& scratch, bigInt& out);
        friend class x509cert;
        friend class tls;
        friend class rsa_pss_sha256;
        friend class rsa_pss_sha384;
        friend class rsa_pkcs1_sha256;
    };
};
