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

// IntFactorization.cpp
// Implementation of the complete prime factorization dispatcher

#include "math/core/mp/Int/IntFactorization.hpp"
#include "math/core/mp/Int/IntOps.hpp"
#include "math/core/mp/Int/IntSqrt.hpp"
#include "math/core/mp/ThreadPool.hpp"

#include <map>
#include <cstdint>
#include <atomic>
#include <mutex>
#include <future>
#include <cmath>

namespace sangi {

// ============================================================================
// uint64-specific prime factorization — fast path that never uses Int objects
// ============================================================================

namespace {

// Smallest prime factor table (SPF) — odd numbers only, uint16 (SPF of a composite ≤ √TABLE_LIMIT)
// TABLE_LIMIT = 1,000,000 → √ ≈ 1000 → uint16 is sufficient
// Memory: 500,000 × 2 bytes = 1 MB
static constexpr int TABLE_LIMIT = 1000000;

struct SpfTable {
    uint16_t data[(TABLE_LIMIT + 1) / 2];  // data[n/2] = SPF of odd n (0 = prime)

    SpfTable() {
        // Initialize all entries to 0 (prime)
        std::memset(data, 0, sizeof(data));
        data[0] = 1;  // n=1 is special

        int sqrt_limit = static_cast<int>(std::sqrt(static_cast<double>(TABLE_LIMIT)));
        for (int p = 3; p <= sqrt_limit; p += 2) {
            if (data[p >> 1] == 0) {  // p is prime
                // Record the SPF for odd multiples of p (only if unset → smallest prime factor)
                for (int64_t m = static_cast<int64_t>(p) * p; m <= TABLE_LIMIT; m += 2 * p) {
                    if (data[m >> 1] == 0) {
                        data[m >> 1] = static_cast<uint16_t>(p);
                    }
                }
            }
        }
    }
};

const SpfTable& getSpfTable() {
    static const SpfTable table;
    return table;
}

// Build the table at process startup (avoids the latency on the first factorize() call)
static const auto& spf_init_ = getSpfTable();

// Forward declaration of pollardRho64 (already defined in IntPrime.cpp's anonymous namespace, but cannot be linked)
uint64_t mulmod64(uint64_t a, uint64_t b, uint64_t m) {
#ifdef _MSC_VER
    uint64_t hi;
    uint64_t lo = _umul128(a, b, &hi);
    uint64_t rem;
    _udiv128(hi, lo, m, &rem);
    return rem;
#else
    return static_cast<uint64_t>(
        (static_cast<unsigned __int128>(a) * b) % m);
#endif
}

uint64_t gcd64(uint64_t a, uint64_t b) {
    if (a == 0) return b;
    if (b == 0) return a;
#ifdef _MSC_VER
    unsigned long shift;
    _BitScanForward64(&shift, a | b);
    unsigned long sa; _BitScanForward64(&sa, a); a >>= sa;
    unsigned long sb; _BitScanForward64(&sb, b); b >>= sb;
#else
    int shift = __builtin_ctzll(a | b);
    a >>= __builtin_ctzll(a);
    b >>= __builtin_ctzll(b);
#endif
    while (true) {
        if (a > b) { uint64_t t = a; a = b; b = t; }
        b -= a;
        if (b == 0) return a << shift;
#ifdef _MSC_VER
        unsigned long ctz; _BitScanForward64(&ctz, b); b >>= ctz;
#else
        b >>= __builtin_ctzll(b);
#endif
    }
}

uint64_t pollardRho64_local(uint64_t n) {
    if (n % 2 == 0) return 2;
    if (n < 4) return n;

    for (uint64_t c = 1; c <= 20; ++c) {
        uint64_t x = 2, y = 2, d = 1;
        auto f = [&](uint64_t v) { return (mulmod64(v, v, n) + c) % n; };

        uint64_t r = 1, q = 1;
        while (d == 1) {
            x = y;
            for (uint64_t i = 0; i < r; ++i) y = f(y);

            uint64_t k = 0;
            while (k < r && d == 1) {
                uint64_t ys = y;
                uint64_t m = (r - k < 128) ? (r - k) : 128;
                for (uint64_t i = 0; i < m; ++i) {
                    y = f(y);
                    uint64_t diff = (x > y) ? (x - y) : (y - x);
                    q = mulmod64(q, diff, n);
                }
                d = gcd64(q, n);
                k += m;

                if (d == n) {
                    // Backtrack
                    y = ys;
                    d = 1;
                    while (d == 1) {
                        y = f(y);
                        uint64_t diff = (x > y) ? (x - y) : (y - x);
                        d = gcd64(diff, n);
                    }
                }
            }
            r <<= 1;
            q = 1;
            if (r > 1000000) break;
        }
        if (d != 1 && d != n) return d;
    }
    return n;  // Not found
}

// Deterministic Miller-Rabin (uint64 range)
// Deterministic below 2^64 with bases {2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37}
bool isPrime64(uint64_t n) {
    if (n < 2) return false;
    if (n < 4) return true;
    if (n % 2 == 0) return false;
    if (n % 3 == 0) return n == 3;
    if (n % 5 == 0) return n == 5;
    if (n % 7 == 0) return n == 7;

    // n-1 = d * 2^r
    uint64_t d = n - 1;
    int r = 0;
    while ((d & 1) == 0) { d >>= 1; ++r; }

    // Modular exponentiation (binary method)
    auto powmod = [](uint64_t base, uint64_t exp, uint64_t mod) -> uint64_t {
        uint64_t result = 1;
        base %= mod;
        while (exp > 0) {
            if (exp & 1) result = mulmod64(result, base, mod);
            exp >>= 1;
            base = mulmod64(base, base, mod);
        }
        return result;
    };

    // For n < 3,215,031,751, bases {2, 3, 5, 7} are sufficient
    // For n < 2^64, {2,3,5,7,11,13,17,19,23,29,31,37} are deterministic
    static const uint64_t witnesses[] = {2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37};
    int num_witnesses = (n < 3215031751ULL) ? 4 : 12;

    for (int i = 0; i < num_witnesses; ++i) {
        uint64_t a = witnesses[i];
        if (a >= n) continue;

        uint64_t x = powmod(a, d, n);
        if (x == 1 || x == n - 1) continue;

        bool composite = true;
        for (int j = 0; j < r - 1; ++j) {
            x = mulmod64(x, x, n);
            if (x == n - 1) { composite = false; break; }
        }
        if (composite) return false;
    }
    return true;
}

// Helper to add a prime factor to factors
void addFactor64(std::vector<std::pair<Int, int>>& factors, uint64_t p, int exp) {
    Int pi(static_cast<int64_t>(p));
    for (auto& [fp, fe] : factors) {
        if (fp == pi) { fe += exp; return; }
    }
    factors.push_back({std::move(pi), exp});
}

// Factorize using the SPF table (n ≤ TABLE_LIMIT, n is odd)
void factorize_spf(uint64_t n, std::vector<std::pair<Int, int>>& factors) {
    const auto& spf = getSpfTable();
    while (n > 1) {
        if ((n & 1) == 0) {
            int exp = 0;
            while ((n & 1) == 0) { n >>= 1; ++exp; }
            addFactor64(factors, 2, exp);
            continue;
        }
        uint16_t s = spf.data[n >> 1];
        if (s == 0) {
            // n is prime
            addFactor64(factors, n, 1);
            return;
        }
        uint64_t p = s;
        int exp = 0;
        while (n % p == 0) { n /= p; ++exp; }
        addFactor64(factors, p, exp);
    }
}

// uint64 Fermat method — instantly finds balanced semiprimes
uint64_t fermat64(uint64_t n, int max_iter = 1000) {
    // ceil(sqrt(n)) — the double approximation is sufficient (error ≤ 1 in the 64-bit range)
    uint64_t a = static_cast<uint64_t>(std::ceil(std::sqrt(static_cast<double>(n))));
    // Correct for the precision shortfall of double
    while (a * a < n) ++a;

    for (int i = 0; i < max_iter; ++i) {
        uint64_t a2 = a * a;
        uint64_t b2 = a2 - n;
        uint64_t b = static_cast<uint64_t>(std::sqrt(static_cast<double>(b2)));
        // Correct the double approximation
        while (b * b < b2) ++b;
        if (b * b == b2) {
            uint64_t f = a - b;
            if (f > 1 && f < n) return f;
        }
        ++a;
    }
    return n;
}

// uint64-specific prime factorization — fast path that never uses Int objects
void factorize_u64(uint64_t n, std::vector<std::pair<Int, int>>& factors) {
    if (n <= 1) return;

    // Handle even numbers
    if ((n & 1) == 0) {
        int exp = 0;
        while ((n & 1) == 0) { n >>= 1; ++exp; }
        addFactor64(factors, 2, exp);
        if (n <= 1) return;
    }

    // Within the table range: factorize instantly via an SPF table lookup
    if (n <= TABLE_LIMIT) {
        factorize_spf(n, factors);
        return;
    }

    // Outside the table range (n > 10^6):
    // 1. The SPF table sieve already covers factors below 1000 → only try 1000 and above
    //    However, since n itself came from outside the table, small factors must also be tried
    //    → the SPF table's prime list cannot be used, so run trial division over 3~997

    // Trial division by primes 3~997 (the range sieved below 1000 when building the SPF table)
    // Primes can be obtained from the SPF table: prime if data[p/2] == 0
    {
        const auto& spf = getSpfTable();
        for (uint64_t p = 3; p < 1000 && p * p <= n; p += 2) {
            if (spf.data[p >> 1] != 0) continue;  // Skip composites
            if (n % p == 0) {
                int exp = 0;
                while (n % p == 0) { n /= p; ++exp; }
                addFactor64(factors, p, exp);
            }
        }
        if (n <= 1) return;
        // If n now fits within the table range
        if (n <= TABLE_LIMIT) { factorize_spf(n, factors); return; }
    }

    // Trial division over 1009~1999 (beyond the SPF table's √TABLE_LIMIT ≈ 1000)
    for (uint64_t p = 1009; p < 2000 && p * p <= n; p += 2) {
        if (n % p == 0) {
            int exp = 0;
            while (n % p == 0) { n /= p; ++exp; }
            addFactor64(factors, p, exp);
        }
    }
    if (n <= 1) return;
    if (n <= TABLE_LIMIT) { factorize_spf(n, factors); return; }

    // If what remains is prime, add it
    if (n < static_cast<uint64_t>(1987) * 1987) {
        // 1987² ≈ less than 3.95 million → below the square of the trial-division range, so definitely prime
        addFactor64(factors, n, 1);
        return;
    }
    if (isPrime64(n)) {
        addFactor64(factors, n, 1);
        return;
    }

    // Composite: split with Fermat + pollardRho64 and recurse
    std::vector<uint64_t> composites;
    composites.push_back(n);

    while (!composites.empty()) {
        uint64_t c = composites.back();
        composites.pop_back();

        if (c <= 1) continue;

        // Factorize via SPF if within the table range
        if (c <= TABLE_LIMIT) {
            factorize_spf(c, factors);
            continue;
        }

        if (isPrime64(c)) {
            addFactor64(factors, c, 1);
            continue;
        }

        // Fermat method (instantly finds balanced semiprimes)
        uint64_t f = fermat64(c, 1000);
        if (f > 1 && f < c) {
            composites.push_back(f);
            composites.push_back(c / f);
            continue;
        }

        // Pollard rho
        f = pollardRho64_local(c);
        if (f > 1 && f < c) {
            composites.push_back(f);
            composites.push_back(c / f);
            continue;
        }

        // Not found (theoretically unreachable)
        addFactor64(factors, c, 1);
    }
}

} // anonymous namespace

// ============================================================================
// findFactor — single factor search
// Staged application of multiple algorithms
// ============================================================================

Int IntFactorization::findFactor(const Int& n, int millerRabinK) {
    // Even
    if (n.isEven()) return Int(2);

    // 1 or less
    if (n <= 1) return n;

    // Primality check — a deterministic test for small n (few rounds suffice)
    int bits = static_cast<int>(n.bitLength());
    {
        // n < 2^64: k=2 suffices (SPRP-2,3 is deterministic below 2^64)
        // n < 2^128: k=5 suffices (a practically sufficient probability)
        int k = (bits <= 64) ? 2 : (bits <= 128) ? 5 : millerRabinK;
        if (IntPrime::isProbablePrime(n, k)) return n;
    }

    // ================================================================
    // Stage 1: Fermat method — a light initial attempt (< 1us)
    // Instantly finds balanced semiprimes with p ≈ q. Always tried first since the cost is extremely low.
    // 1000 iterations cover the range |p-q| ≤ ~2000·√2 ≈ 2828.
    // ================================================================
    {
        Int f = IntPrime::findFactor_Fermat(n, 1000);
        if (f > Int(1) && f < n) return f;
    }

    // ================================================================
    // Stage 2: Pollard rho / Brent (a few us~a few ms)
    // The most general-purpose method. O(p^{1/2}), so faster for smaller factors.
    // Accelerated with dedicated uint64 (≤63 bits) / 128-bit Montgomery (≤127 bits) versions.
    // ================================================================
    {
        Int f = IntPrime::findFactor_PollardRho(n);
        if (f > Int(1) && f < n) return f;
    }

    // ================================================================
    // Stage 3: special methods — selected by digit count
    // ================================================================

    // 3a: SQUFOF (≤60 bits only, O(n^{1/4}), minimal memory)
    // A small-to-medium-scale fallback when rho did not find a factor.
    if (bits <= 60) {
        Int f = IntPrime::findFactor_SQUFOF(n);
        if (f > Int(1) && f < n) return f;
    }

    // 3b: Pollard p-1 — run only when bits >= 50
    // For small numbers rho is sufficient. B1 is tuned according to bits.
    if (bits >= 50) {
        uint64_t B1 = (bits <= 80) ? 5000 : 50000;
        Int f = IntPrime::findFactor_PollardP1(n, B1);
        if (f > Int(1) && f < n) return f;
    }

    // 3c: Williams p+1 — run only when bits >= 50
    if (bits >= 50) {
        uint64_t B1 = (bits <= 80) ? 5000 : 50000;
        Int f = IntPrime::findFactor_WilliamsP1(n, B1);
        if (f > Int(1) && f < n) return f;
    }

    // ================================================================
    // Stage 4: heavy factorization
    // ================================================================
    // ECM factorization (sequential)
    {
        int curves = (bits <= 80) ? 50 : (bits <= 120) ? 100 : 200;
        Int f = IntPrime::findFactor_ECM(n, curves, 0);
        if (f > Int(1) && f < n) return f;
    }

    return n;  // No factor was found
}

// ============================================================================
// factorize — complete prime factorization
// ============================================================================

std::vector<std::pair<Int, int>> IntFactorization::factorize(const Int& n, int millerRabinK) {
    // Special cases
    if (n.isNaN() || n.isInfinite()) return {};
    if (n.isZero()) return {{Int(0), 1}};

    Int remaining = n;
    if (remaining.isNegative()) remaining = -remaining;
    if (remaining.isOne()) return {};

    std::vector<std::pair<Int, int>> factors;

    // Stage 1: remove the prime factor 2 (fast via bit shifts)
    if (remaining.isEven()) {
        int exp2 = 0;
        while (remaining.isEven()) {
            remaining >>= 1;
            ++exp2;
        }
        factors.push_back({Int(2), exp2});
    }

    if (remaining.isOne()) {
        std::sort(factors.begin(), factors.end());
        return factors;
    }

    // ================================================================
    // Stage 1.5: uint64 fast path — for 1 limb, factorize without ever using Int
    // Completed with the SPF table (≤ 10^6) + trial division + pollardRho64.
    // Every number of 19 digits or fewer is handled by this path.
    // ================================================================
    if (remaining.bitLength() <= 63) {
        uint64_t n64 = remaining.toUInt64();
        factorize_u64(n64, factors);
        std::sort(factors.begin(), factors.end());
        return factors;
    }

    // Stage 2: remove factors in 3~1987 via small-prime batch GCD
    IntPrime::extractSmallPrimeFactors(remaining, factors);

    // Done if remaining == 1
    if (remaining.isOne()) {
        std::sort(factors.begin(), factors.end());
        return factors;
    }

    // Stage 3: if the remainder is prime, add it and finish
    if (IntPrime::isProbablePrime(remaining, millerRabinK)) {
        factors.push_back({remaining, 1});
        std::sort(factors.begin(), factors.end());
        return factors;
    }

    // Stage 4: recursively factorize the remainder
    // Manage composites with a stack
    std::vector<Int> composites;
    composites.push_back(remaining);

    while (!composites.empty()) {
        Int c = composites.back();
        composites.pop_back();

        if (c.isOne()) continue;

        if (IntPrime::isProbablePrime(c, millerRabinK)) {
            // Prime: add to the factor list
            bool found = false;
            for (auto& [p, e] : factors) {
                if (p == c) {
                    ++e;
                    found = true;
                    break;
                }
            }
            if (!found) factors.push_back({c, 1});
            continue;
        }

        // Composite: find one factor and split
        Int f = findFactor(c, millerRabinK);
        if (f == c || f <= 1) {
            // No factor was found (theoretically unreachable, but a safety measure)
            // Add the remainder itself as a factor
            bool found = false;
            for (auto& [p, e] : factors) {
                if (p == c) {
                    ++e;
                    found = true;
                    break;
                }
            }
            if (!found) factors.push_back({c, 1});
            continue;
        }

        Int cofactor = c / f;

        // Push f and cofactor onto the stack and process recursively
        composites.push_back(f);
        composites.push_back(cofactor);
    }

    // Sort in ascending order
    std::sort(factors.begin(), factors.end());
    return factors;
}

// ============================================================================
// toString — string representation of the factorization result
// ============================================================================

std::string IntFactorization::toString(const std::vector<std::pair<Int, int>>& factors) {
    if (factors.empty()) return "1";

    std::ostringstream oss;
    bool first = true;
    for (auto& [p, e] : factors) {
        if (!first) oss << " * ";
        first = false;
        oss << p.toString();
        if (e > 1) oss << "^" << e;
    }
    return oss.str();
}

} // namespace sangi