// Copyright (c) 2025 The Limenka developers
// Distributed under the MIT software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.

#ifndef LIMENKA_CRYPTO_BIGNUM_H
#define LIMENKA_CRYPTO_BIGNUM_H

#include <cstddef>
#include <cstdint>
#include <memory>
#include <span>
#include <string>
#include <vector>

class BigNum {
public:
    BigNum() { m_limbs = {0}; }
    explicit BigNum(const std::vector<uint8_t>& bytes, bool big_endian = true);
    explicit BigNum(uint32_t val);

    std::vector<uint8_t> to_bytes(size_t width = 0) const;

    bool is_zero() const;
    bool is_one() const;
    bool is_even() const { return (m_limbs[0] & 1) == 0; }
    size_t num_limbs() const { return m_limbs.size(); }
    uint32_t limb(size_t i) const { return i < m_limbs.size() ? m_limbs[i] : 0; }
    size_t bit_length() const;

    int compare(const BigNum& other) const;
    bool operator==(const BigNum& other) const { return compare(other) == 0; }
    bool operator!=(const BigNum& other) const { return compare(other) != 0; }

    BigNum& operator+=(const BigNum& other);
    BigNum& operator-=(const BigNum& other);
    BigNum& operator*=(const BigNum& other);

    BigNum operator+(const BigNum& other) const { BigNum r(*this); r += other; return r; }
    BigNum operator-(const BigNum& other) const { BigNum r(*this); r -= other; return r; }
    BigNum operator*(const BigNum& other) const { BigNum r(*this); r *= other; return r; }
    BigNum operator%(const BigNum& other) const { BigNum r(*this); r %= other; return r; }

    BigNum& operator%=(const BigNum& mod);

    /** Sequential squaring: sqr = x, then sqr = sqr*sqr mod m, D times.
     *  Stores every k-th intermediate (k=16) for proof reconstruction.
     *  intermediates[i] = x^(2^(i*k)) mod m.  intermediates[0] = x^(2^0) = x.
     */
    static BigNum sqr_chain(const BigNum& base, uint32_t D, const BigNum& m,
                            std::vector<BigNum>* intermediates = nullptr);

    /** Modular exponentiation: base^exp mod m (square-and-multiply). */
    static BigNum mod_pow(BigNum base, BigNum exp, const BigNum& m);

    /** Montgomery modular multiplication: (a * b * R^-1) mod m.  m must be odd. */
    static BigNum mont_mul(const BigNum& a, const BigNum& b,
                           const BigNum& m, uint32_t m_prime);

    /** Montgomery precomputation: m_prime such that m * m_prime ≡ -1 (mod 2^32). */
    static uint32_t mont_mp(const BigNum& m);

    /** Montgomery precomputation: R^2 mod m where R = 2^(32 * m.num_limbs()). */
    static BigNum mont_r2(const BigNum& m);

    /** Binary GCD. */
    static BigNum gcd(const BigNum& a, const BigNum& b);

    /** Division: compute quotient and remainder.  num / den = q, num % den = r. */
    static void div_rem(const BigNum& num, const BigNum& den,
                        BigNum& quot, BigNum& rem);

    // Allow hash functions to access internal representation
    const std::vector<uint32_t>& limbs() const { return m_limbs; }

private:
    std::vector<uint32_t> m_limbs;
    void trim();
};

std::string BigNumToHex(const BigNum& n);

#endif // LIMENKA_CRYPTO_BIGNUM_H