// Copyright (C) 2026 Kiyotsugu Arai
// SPDX-License-Identifier: LGPL-3.0-or-later

// IntOps.hpp
// Utilities for operations on arbitrary-precision integers

#ifndef SANGI_INT_OPS_HPP
#define SANGI_INT_OPS_HPP

#include <math/core/mp/Int/IntBase.hpp>
#include <math/core/mp/Int/MpnOps.hpp>

namespace sangi {

    /**
     * @brief Utility class for operations on arbitrary-precision integers
     *
     * This class provides static methods that operate on the internal value of
     * the Int class. It mainly contains implementations of arithmetic and
     * bitwise operations.
     */
    class SANGI_API IntOps {
    public:
        // Comparison operations
        static bool compareAbsLess(const Int& lhs, const Int& rhs);
        static bool compareAbsGreater(const Int& lhs, const Int& rhs);
        static bool compareAbsEqual(const Int& lhs, const Int& rhs);

        // Addition / subtraction operations
        static void addAbsolute(Int& result, const Int& other);
        static void addAbsolute(const Int& lhs, const Int& rhs, Int& result);
        static void subtractAbsolute(const Int& lhs, const Int& rhs, Int& result);
        static void subtractAbsoluteInPlace(Int& result, const Int& other);

        // Increment / decrement: avoid Int(1) construction + operator+/-
        // delta = +1 or -1
        static void addDelta(Int& value, int delta);

        // Three-argument add/sub/mul/div: reuse result's buffer (GMP mpz_add/sub/mul/tdiv_q equivalents)
        // result may alias lhs or rhs
        static void add(const Int& lhs, const Int& rhs, Int& result);
        static void sub(const Int& lhs, const Int& rhs, Int& result);
        static void mul(const Int& lhs, const Int& rhs, Int& result);
        static void div(const Int& dividend, const Int& divisor, Int& result);

        // Unchecked versions: skip isSpecialState() (caller guarantees the precondition)
        // Used on the hot paths of internal algorithms
        static void addUnchecked(const Int& lhs, const Int& rhs, Int& result);
        static void subUnchecked(const Int& lhs, const Int& rhs, Int& result);
        static void mulUnchecked(const Int& lhs, const Int& rhs, Int& result);
        static void divUnchecked(const Int& dividend, const Int& divisor, Int& result);

        // Single-word operations (equivalent to GMP _ui: avoid constructing a temporary Int)
        // addWord: result += word (signed addition)
        static void addWord(Int& result, uint64_t word);
        // subWord: result -= word (signed subtraction)
        static void subWord(Int& result, uint64_t word);
        // divWord: result /= word, returns the remainder
        static uint64_t divWord(Int& result, uint64_t word);
        // divExactWord: result /= word (precondition: word > 0 and word divides result)
        // Uses divexact_by_odd (Hensel inverse); since no quotient estimation is needed,
        // it is about 2x faster than divmod_1.
        // Behavior is undefined if word does not divide result. Used in num/g and den/g
        // inside Rational::reduce().
        static void divExactWord(Int& result, uint64_t word);

        // Multiplication operations
        static void multiplyAbsolute(Int& result, const Int& other);
        static void multiplyAbsolute(const Int& lhs, const Int& rhs, Int& result);
        static void multiplyWord(Int& result, uint64_t word);
        static void square(const Int& value, Int& result);

        // Unchecked versions: skip isSpecialState() (caller guarantees the precondition)
        static void squareUnchecked(const Int& value, Int& result);
        static void multiplyWordUnchecked(Int& result, uint64_t word);

        // NTT-cached absolute multiplication: cache the forward NTT of rhs for reuse.
        // Effective when multiplying the same rhs by multiple lhs values (e.g. QR
        // reuse in BS merge).
        static void mulAbsCached(const Int& lhs, const Int& rhs, Int& result,
                                 prime_ntt::NttCache& cache);

        // YC-2: Fused multiply-add: result = a*b + c*d (accelerated via NTT fusion)
        // Computes the T = TL*QR + PL*TR pattern from BS merge in a single inverse NTT + CRT.
        // Handles signs (positive+positive, positive-negative, etc.).
        static void mulAdd(const Int& a, const Int& b,
                           const Int& c, const Int& d, Int& result);

        // Division and remainder operations
        static void divideAbsolute(Int& result, const Int& divisor);
        static void moduloAbsolute(Int& result, const Int& divisor);
        static Int divmod(const Int& dividend, const Int& divisor, Int& remainder);

        // Unchecked versions: used to remove duplicate checks from operator/ and operator%
        // Preconditions: result and divisor are non-special, divisor != 0
        static void divideAbsoluteUnchecked(Int& result, const Int& divisor);
        static void moduloAbsoluteUnchecked(Int& result, const Int& divisor);

        // Bitwise operations
        static void bitwiseAnd(Int& result, const Int& other);
        static void bitwiseOr(Int& result, const Int& other);
        static void bitwiseXor(Int& result, const Int& other);
        static void bitwiseNot(Int& result);
        static void leftShift(Int& value, int shift);
        static void rightShift(Int& value, int shift);
        // Three-argument versions: result = value << shift / value >> shift (buffer reuse)
        static void leftShift(const Int& value, int shift, Int& result);
        static void rightShift(const Int& value, int shift, Int& result);

        // Unchecked versions: avoid duplicate checks
        static void bitwiseAndUnchecked(Int& result, const Int& other);
        static void bitwiseOrUnchecked(Int& result, const Int& other);
        static void bitwiseXorUnchecked(Int& result, const Int& other);
        static void bitwiseNotUnchecked(Int& result);
        // Three-argument shift, unchecked: for inner loops of powerMod, etc.
        static void leftShiftUnchecked(const Int& value, int shift, Int& result);
        static void rightShiftUnchecked(const Int& value, int shift, Int& result);

        // Coefficient of power of 2
        static Int pow2(uint64_t exponent);

        // GCD (greatest common divisor)
        static Int gcd(const Int& a, const Int& b);

        // Factor removal (removeFactor)
        // Remove all occurrences of factor from value and return how many were removed.
        // Satisfies value = result * factor^count.
        // GMP compatible: mpz_remove(rop, op, f)
        // Returns: number of times the factor was removed (exponent)
        // result: value after removal (value / factor^count)
        static uint64_t removeFactor(const Int& value, const Int& factor, Int& result);

        // Fused multiply-add / multiply-subtract: rop += a*b / rop -= a*b
        // (equivalent to GMP mpz_addmul / mpz_submul)
        // Minimizes the construction of intermediate Int temporaries.
        static void addmul(Int& rop, const Int& a, const Int& b);
        static void submul(Int& rop, const Int& a, const Int& b);

        // Single-word variants: rop += a*word / rop -= a*word
        static void addmul(Int& rop, const Int& a, uint64_t word);
        static void submul(Int& rop, const Int& a, uint64_t word);

        // Square root
        static Int sqrt(const Int& value);

        // Floor division
        // Division rounded toward negative infinity.
        // Example: floorDiv(-7, 3) = -3 (ordinary division gives -2)
        static Int floorDiv(const Int& dividend, const Int& divisor);

        // Floor modulo
        // The sign of the result is always the same as the divisor (same as Python's %).
        // Example: floorMod(-7, 3) = 2, floorMod(7, -3) = -2
        static Int floorMod(const Int& dividend, const Int& divisor);

        // Return floor division and modulo together
        // quotient = floorDiv(a, b), remainder = floorMod(a, b)
        static Int floorDivMod(const Int& dividend, const Int& divisor, Int& remainder);

        // Ceiling division
        // Division rounded toward positive infinity.
        // Example: ceilDiv(7, 3) = 3, ceilDiv(-7, 3) = -2
        static Int ceilDiv(const Int& dividend, const Int& divisor);

        // Exact division (when the divisor is guaranteed to divide the dividend).
        // Behavior is undefined when not divisible (the remainder check is skipped for speed).
        // Multi-limb divisor: uses mpn::divexact (Hensel lifting), lighter than normal division.
        static Int divExact(const Int& dividend, const Int& divisor);

        // divExactInPlace: result /= divisor (precondition: divisor divides result)
        // Chooses between divExactWord (1-limb) and mpn::divexact (multi-limb).
        // Used in the multi-limb gcd path of Rational::reduce().
        static void divExactInPlace(Int& result, const Int& divisor);

    private:
        // Multiplication via the Karatsuba algorithm
        static void karatsubaMultiply(const Int& a, const Int& b, Int& result);

        // Multiplication via Toom-Cook-3
        static void toomCookMultiply(const Int& a, const Int& b, Int& result);

        // Multiplication via Toom-Cook-4
        static void toomCook4Multiply(const Int& a, const Int& b, Int& result);

        // General multiplication (handles unbalanced inputs, automatic selection by size)
        static void generalMultiply(const Int& a, const Int& b, Int& result);

        // Multiplication via FFT
        static void fftMultiply(const Int& a, const Int& b, Int& result);

        // Division via Newton's method
        static void newtonDivision(const Int& dividend, const Int& divisor, Int& result);
    };

    // ----------------------------------------------------------------
    // Inline implementations: addition / subtraction (three-argument form)
    // Marking these inline allows operator+/- to call straight through to
    // mpn::add/sub without a function-call indirection, eliminating overhead
    // for small operations.
    // ----------------------------------------------------------------

    inline void IntOps::addAbsolute(const Int& lhs, const Int& rhs, Int& result) {
        size_t lhsSize = lhs.m_words.size();
        size_t rhsSize = rhs.m_words.size();

        // mpn::add requires an >= bn.
        // If result aliases lhs or rhs, perform in-place addition reusing the
        // resized buffer (to avoid an unnecessary copy).
        size_t bigSize = std::max(lhsSize, rhsSize);
        size_t smallSize = std::min(lhsSize, rhsSize);

        if (&result == &lhs || &result == &rhs) {
            // result aliases one of lhs or rhs
            const Int& other = (&result == &lhs) ? rhs : lhs;
            size_t otherSize = other.m_words.size();
            size_t resultSize = result.m_words.size();

            // Extend result to bigSize+1 (existing data is preserved)
            result.m_words.resize_uninitialized(bigSize + 1);
            uint64_t* rp = result.m_words.data();
            const uint64_t* op = other.m_words.data();

            uint64_t carry;
            if (resultSize >= otherSize) {
                // result is the larger side: r == a aliasing
                carry = mpn::add(rp, rp, resultSize, op, otherSize);
            } else {
                // result is the smaller side: r == b aliasing
                // Safe because mpn::add reads b[i] before writing r[i] for i < bn.
                carry = mpn::add(rp, op, otherSize, rp, resultSize);
            }
            if (carry) {
                rp[bigSize] = 1;
            } else {
                result.m_words.resize_uninitialized(bigSize);
            }
        } else {
            // result is a separate object from both lhs and rhs
            result.m_words.resize_uninitialized(bigSize + 1);
            const uint64_t* bigData = (lhsSize >= rhsSize) ? lhs.m_words.data() : rhs.m_words.data();
            const uint64_t* smallData = (lhsSize >= rhsSize) ? rhs.m_words.data() : lhs.m_words.data();
            uint64_t carry = mpn::add(result.m_words.data(),
                                       bigData, bigSize, smallData, smallSize);
            if (carry) {
                result.m_words.data()[bigSize] = 1;
            } else {
                result.m_words.resize_uninitialized(bigSize);
            }
        }
    }

    inline void IntOps::subtractAbsolute(const Int& lhs, const Int& rhs, Int& result) {
        // Precondition: |lhs| >= |rhs| (guaranteed by the caller)
        size_t lhsSize = lhs.m_words.size();
        size_t rhsSize = rhs.m_words.size();

        if (&result == &lhs) {
            // result == lhs: in-place subtraction (r == a aliasing)
            mpn::sub(result.m_words.data(),
                     result.m_words.data(), lhsSize,
                     rhs.m_words.data(), rhsSize);
        } else {
            result.m_words.resize_uninitialized(lhsSize);
            mpn::sub(result.m_words.data(),
                     lhs.m_words.data(), lhsSize,
                     rhs.m_words.data(), rhsSize);
        }

        // Normalize: trim leading zeros
        while (result.m_words.size() > 0 && result.m_words.back() == 0) {
            result.m_words.pop_back();
        }
        if (result.m_words.empty()) {
            result.setSign(0);
        }
    }

    // ----------------------------------------------------------------
    // Three-argument add/sub: reuse result's buffer
    // Equivalent to GMP's mpz_add(r, a, b) / mpz_sub(r, a, b)
    // ----------------------------------------------------------------

    // --- Unchecked version (isSpecialState check skipped) ---

    inline void IntOps::addUnchecked(const Int& lhs, const Int& rhs, Int& result) {
        if (lhs.m_sign == 0) { if (&result != &rhs) result = rhs; return; }
        if (rhs.m_sign == 0) { if (&result != &lhs) result = lhs; return; }

        if (lhs.m_sign == rhs.m_sign) {
            addAbsolute(lhs, rhs, result);
            result.m_sign = lhs.m_sign;
        } else {
            int cmp = mpn::cmp(lhs.m_words.data(), lhs.m_words.size(),
                               rhs.m_words.data(), rhs.m_words.size());
            if (cmp < 0) {
                subtractAbsolute(rhs, lhs, result);
                result.m_sign = rhs.m_sign;
            } else if (cmp > 0) {
                subtractAbsolute(lhs, rhs, result);
                result.m_sign = lhs.m_sign;
            } else {
                result.m_words.clear();
                result.m_sign = 0;
            }
        }
        result.m_state = NumericState::Normal;
    }

    inline void IntOps::subUnchecked(const Int& lhs, const Int& rhs, Int& result) {
        if (rhs.m_sign == 0) { if (&result != &lhs) result = lhs; return; }
        if (lhs.m_sign == 0) {
            if (&result != &rhs) result = rhs;
            result.m_sign = -result.m_sign;
            return;
        }
        if (lhs.m_sign != rhs.m_sign) {
            addAbsolute(lhs, rhs, result);
            result.m_sign = lhs.m_sign;
        } else {
            int cmp = mpn::cmp(lhs.m_words.data(), lhs.m_words.size(),
                               rhs.m_words.data(), rhs.m_words.size());
            if (cmp >= 0) {
                subtractAbsolute(lhs, rhs, result);
                result.m_sign = (cmp == 0) ? 0 : lhs.m_sign;
            } else {
                subtractAbsolute(rhs, lhs, result);
                result.m_sign = -rhs.m_sign;
            }
        }
        result.m_state = NumericState::Normal;
    }

    // --- Checked version (thin wrapper) ---

    inline void IntOps::add(const Int& lhs, const Int& rhs, Int& result) {
        if (lhs.isSpecialState() || rhs.isSpecialState()) [[unlikely]] {
            result = lhs + rhs;
            return;
        }
        addUnchecked(lhs, rhs, result);
    }

    inline void IntOps::sub(const Int& lhs, const Int& rhs, Int& result) {
        if (lhs.isSpecialState() || rhs.isSpecialState()) [[unlikely]] {
            result = lhs - rhs;
            return;
        }
        subUnchecked(lhs, rhs, result);
    }

} // namespace sangi

#endif // SANGI_INT_OPS_HPP