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

// IntGCD.cpp
// Implementation of GCD and LCM operations for Int
//
// GCD: Lehmer + recursive HGCD (O(M(n) log n))
// extendedGcd: Lehmer + HGCD + cofactor tracking (O(M(n) log² n))
//   - Apply the HGCD matrix to the cofactor pair for speedup
//   - Track only one cofactor (x), recover y = (g - a*x) / b

#include "math/core/mp/Int/IntGCD.hpp"
#include "math/core/mp/Int/IntOps.hpp"
#include "math/core/mp/Int/IntSpecialStates.hpp"
#include "math/core/mp/Int/MpnOps.hpp"
#include "math/core/mp/Int/ScratchArena.hpp"
#include <algorithm>
#include <bit>

namespace sangi {

// =========================================================================
// GCD (existing: unchanged)
// =========================================================================

Int IntGCD::gcd(const Int& a, const Int& b) {
    // 1. Handle special states
    if (a.isNaN() || b.isNaN()) [[unlikely]] {
        return Int::NaN();
    }
    if (a.isInfinite() || b.isInfinite()) [[unlikely]] {
        return Int::NaN();
    }

    // 2. Handle zero cases
    if (a.isZero() && b.isZero()) [[unlikely]] {
        return Int::Zero();
    }
    if (a.isZero()) [[unlikely]] {
        return abs(b);
    }
    if (b.isZero()) [[unlikely]] {
        return abs(a);
    }

    size_t an = a.size();
    size_t bn = b.size();

    // =========================================================
    // 3. Fast path: 1-2 limb inputs (avoid arena/vector allocation)
    // =========================================================
    if (an <= 2 && bn <= 2) {
        uint64_t u0 = a.word(0), u1 = a.word(1);
        uint64_t v0 = b.word(0), v1 = b.word(1);

        size_t u_ctz = u0 ? std::countr_zero(u0) : 64 + std::countr_zero(u1);
        size_t v_ctz = v0 ? std::countr_zero(v0) : 64 + std::countr_zero(v1);
        size_t k = std::min(u_ctz, v_ctz);

        if (u_ctz >= 64) {
            u0 = u1 >> (u_ctz - 64); u1 = 0;
        } else if (u_ctz > 0) {
            u0 = (u0 >> u_ctz) | (u1 << (64 - u_ctz)); u1 >>= u_ctz;
        }
        if (v_ctz >= 64) {
            v0 = v1 >> (v_ctz - 64); v1 = 0;
        } else if (v_ctz > 0) {
            v0 = (v0 >> v_ctz) | (v1 << (64 - v_ctz)); v1 >>= v_ctz;
        }

        uint64_t g0, g1;
        if (u1 == 0 && v1 == 0) {
            g0 = mpn::gcd_11(u0, v0);
            g1 = 0;
        } else {
            mpn::DoubleLimb g = mpn::gcd_22(u1, u0, v1, v0);
            g0 = g.d0;
            g1 = g.d1;
        }

        if (k > 0) {
            if (k >= 64) {
                g1 = g0 << (k - 64);
                g0 = 0;
            } else {
                g1 = (g1 << k) | (g0 >> (64 - k));
                g0 <<= k;
            }
        }

        if (g1 > 0) {
            uint64_t buf[2] = {g0, g1};
            return Int::fromRawWords(std::span<const uint64_t>(buf, 2), +1);
        }
        if (g0 > 0) {
            return Int::fromRawWords(std::span<const uint64_t>(&g0, 1), +1);
        }
        return Int::Zero();
    }

    // =========================================================
    // 4. General path: Lehmer GCD (uses arena, no vector allocation)
    // =========================================================
    ScratchScope scope;
    size_t max_size = std::max(an, bn);

    size_t buf_size = 2 * max_size + 32;
    uint64_t* up = getThreadArena().alloc_limbs(buf_size);
    uint64_t* vp = getThreadArena().alloc_limbs(buf_size);
    uint64_t* gp = getThreadArena().alloc_limbs(buf_size);

    std::memset(up, 0, buf_size * sizeof(uint64_t));
    std::memset(vp, 0, buf_size * sizeof(uint64_t));
    std::memset(gp, 0, buf_size * sizeof(uint64_t));

    if (an >= bn) {
        for (size_t i = 0; i < an; i++) up[i] = a.word(i);
        for (size_t i = 0; i < bn; i++) vp[i] = b.word(i);
    } else {
        for (size_t i = 0; i < bn; i++) up[i] = b.word(i);
        for (size_t i = 0; i < an; i++) vp[i] = a.word(i);
        std::swap(an, bn);
    }

    size_t k_a = mpn::ctz_limb_array(up, an);
    size_t k_b = mpn::ctz_limb_array(vp, bn);
    size_t k = std::min(k_a, k_b);

    if (k_a > 0) an = mpn::rshift(up, up, an, k_a);
    if (k_b > 0) bn = mpn::rshift(vp, vp, bn, k_b);

    // Removing the common power-of-two factor (rshift) can shrink `up` below `vp`.
    //   e.g. gcd(10^80, b) shrinks up from 10^80 (5 limbs) to 5^80 (3 limbs), so an<bn.
    // mpn::gcd_lehmer assumes the first argument is at least the second (usize >= n);
    // when usize<n it reads stale high limbs of the shrunk side and returns a wrong gcd
    // (the root cause of round28 #2). Re-establish an>=bn here.
    if (an < bn) { std::swap(up, vp); std::swap(an, bn); }

    size_t scratch_size = mpn::gcd_lehmer_scratch_size(an, bn);
    uint64_t* scratch = getThreadArena().alloc_limbs(scratch_size);

    size_t rn = mpn::gcd_lehmer(gp, up, an, vp, bn, scratch);

    if (k > 0 && rn > 0) {
        rn = mpn::lshift_arbitrary(gp, gp, rn, k);
    }

    if (rn == 0) {
        return Int::Zero();
    }
    return Int::fromRawWords(std::span<const uint64_t>(gp, rn), +1);
}

// =========================================================================
// LCM (unchanged)
// =========================================================================

Int IntGCD::lcm(const Int& a, const Int& b) {
    if (a.isNaN() || b.isNaN()) [[unlikely]] { return Int::NaN(); }
    if (a.isInfinite() || b.isInfinite()) [[unlikely]] { return Int::NaN(); }
    if (a.isZero() || b.isZero()) [[unlikely]] { return Int::Zero(); }

    Int g = gcd(a, b);
    if (g.isOne()) {
        return abs(a * b);
    }
    return abs(a) * (abs(b) / g);
}

// =========================================================================
// Extended GCD: HGCD-enabled version
// =========================================================================

namespace {

// ---- Signed multi-precision (non-owning pointer into the ScratchArena) ----
struct SignedMpn {
    uint64_t* d;  // magnitude data
    size_t    n;  // number of limbs (0 = zero value)
    int       s;  // +1, -1, or 0
};

// Set to zero
inline void smn_zero(SignedMpn& x) { x.n = 0; x.s = 0; }

// Set to a single-limb value (buf is the start of the buffer)
inline void smn_set_ui(SignedMpn& x, uint64_t val, int sign) {
    if (val == 0) { x.n = 0; x.s = 0; }
    else { x.d[0] = val; x.n = 1; x.s = sign; }
}

// Signed linear combination: result = u * a - v * b (signed)
// u, v are single limbs (unsigned)
// result is written to r_buf (requires capacity max(a.n, b.n) + 2)
// tmp is a work buffer (requires capacity max(a.n, b.n) + 2)
static void smn_linearcomb_1(
    uint64_t* r_buf, size_t* r_size, int* r_sign,
    uint64_t u, const SignedMpn& a,
    uint64_t v, const SignedMpn& b,
    uint64_t* tmp)
{
    // term1 = u * |a|, sign = a.s
    size_t t1n = 0;
    int t1s = 0;
    if (a.n > 0 && u > 0) {
        uint64_t cy = mpn::mul_1(tmp, a.d, a.n, u);
        t1n = a.n;
        if (cy) tmp[t1n++] = cy;
        t1s = a.s;
    }

    // term2 = v * |b|, sign = b.s
    size_t t2n = 0;
    int t2s = 0;
    if (b.n > 0 && v > 0) {
        uint64_t cy = mpn::mul_1(r_buf, b.d, b.n, v);
        t2n = b.n;
        if (cy) r_buf[t2n++] = cy;
        t2s = b.s;
    }

    // result = t1s * |t1| - t2s * |t2| = t1s * |t1| + (-t2s) * |t2|
    int eff_t2s = -t2s;

    if (t1n == 0) {
        // result = eff_t2s * |t2|  (r_buf already has |t2|)
        *r_size = t2n;
        *r_sign = eff_t2s;
        return;
    }
    if (t2n == 0) {
        std::memcpy(r_buf, tmp, t1n * sizeof(uint64_t));
        *r_size = t1n;
        *r_sign = t1s;
        return;
    }

    if (t1s == eff_t2s) {
        // Same sign: add absolute values
        if (t1n >= t2n) {
            uint64_t cy = mpn::add(r_buf, tmp, t1n, r_buf, t2n);
            *r_size = t1n;
            if (cy) r_buf[(*r_size)++] = cy;
        } else {
            uint64_t cy = mpn::add(r_buf, r_buf, t2n, tmp, t1n);
            *r_size = t2n;
            if (cy) r_buf[(*r_size)++] = cy;
        }
        *r_sign = t1s;
    } else {
        // Different sign: subtract absolute values
        int cmp_result;
        if (t1n != t2n) {
            cmp_result = (t1n > t2n) ? 1 : -1;
        } else {
            cmp_result = mpn::cmp(tmp, t1n, r_buf, t2n);
        }

        if (cmp_result == 0) {
            *r_size = 0; *r_sign = 0;
        } else if (cmp_result > 0) {
            // |t1| > |t2|
            mpn::sub(r_buf, tmp, t1n, r_buf, t2n);
            *r_size = mpn::normalized_size(r_buf, t1n);
            *r_sign = t1s;
        } else {
            // |t2| > |t1|
            mpn::sub(r_buf, r_buf, t2n, tmp, t1n);
            *r_size = mpn::normalized_size(r_buf, t2n);
            *r_sign = eff_t2s;
        }
    }
}

// Multi-precision signed linear combination: result = M_row * a - M_col * b (signed)
// M_row, M_col are multi-precision (unsigned, mn limbs)
// result is written to r_buf (requires capacity mn + max(a.n, b.n) + 4)
// tmp1, tmp2 are work buffers (each requires mn + max(a.n, b.n) + 4)
static void smn_linearcomb_mp(
    uint64_t* r_buf, size_t* r_size, int* r_sign,
    const uint64_t* mp_u, size_t mn_u, const SignedMpn& a,
    const uint64_t* mp_v, size_t mn_v, const SignedMpn& b,
    uint64_t* tmp1, uint64_t* tmp2, uint64_t* mul_scratch)
{
    // term1 = mp_u * |a|, sign = a.s
    size_t t1n = 0;
    int t1s = 0;
    size_t un = mpn::normalized_size(mp_u, mn_u);
    size_t an = a.n;
    if (un > 0 && an > 0) {
        t1n = un + an;
        if (un >= an)
            mpn::multiply(tmp1, mp_u, un, a.d, an, mul_scratch);
        else
            mpn::multiply(tmp1, a.d, an, mp_u, un, mul_scratch);
        t1n = mpn::normalized_size(tmp1, t1n);
        t1s = a.s;
    }

    // term2 = mp_v * |b|, sign = b.s
    size_t t2n = 0;
    int t2s = 0;
    size_t vn = mpn::normalized_size(mp_v, mn_v);
    size_t bn = b.n;
    if (vn > 0 && bn > 0) {
        t2n = vn + bn;
        if (vn >= bn)
            mpn::multiply(tmp2, mp_v, vn, b.d, bn, mul_scratch);
        else
            mpn::multiply(tmp2, b.d, bn, mp_v, vn, mul_scratch);
        t2n = mpn::normalized_size(tmp2, t2n);
        t2s = b.s;
    }

    // result = t1s * |t1| + (-t2s) * |t2|
    int eff_t2s = -t2s;

    if (t1n == 0) {
        if (t2n > 0) std::memcpy(r_buf, tmp2, t2n * sizeof(uint64_t));
        *r_size = t2n; *r_sign = eff_t2s; return;
    }
    if (t2n == 0) {
        std::memcpy(r_buf, tmp1, t1n * sizeof(uint64_t));
        *r_size = t1n; *r_sign = t1s; return;
    }

    if (t1s == eff_t2s) {
        if (t1n >= t2n) {
            std::memcpy(r_buf, tmp1, t1n * sizeof(uint64_t));
            uint64_t cy = mpn::add(r_buf, r_buf, t1n, tmp2, t2n);
            *r_size = t1n;
            if (cy) r_buf[(*r_size)++] = cy;
        } else {
            std::memcpy(r_buf, tmp2, t2n * sizeof(uint64_t));
            uint64_t cy = mpn::add(r_buf, r_buf, t2n, tmp1, t1n);
            *r_size = t2n;
            if (cy) r_buf[(*r_size)++] = cy;
        }
        *r_sign = t1s;
    } else {
        int cmp_result;
        if (t1n != t2n) cmp_result = (t1n > t2n) ? 1 : -1;
        else cmp_result = mpn::cmp(tmp1, t1n, tmp2, t2n);

        if (cmp_result == 0) {
            *r_size = 0; *r_sign = 0;
        } else if (cmp_result > 0) {
            mpn::sub(r_buf, tmp1, t1n, tmp2, t2n);
            *r_size = mpn::normalized_size(r_buf, t1n);
            *r_sign = t1s;
        } else {
            mpn::sub(r_buf, tmp2, t2n, tmp1, t1n);
            *r_size = mpn::normalized_size(r_buf, t2n);
            *r_sign = eff_t2s;
        }
    }
}

// Apply an hgcd2 matrix to the cofactor pair (su, sv)
// new su = M.u[1][1] * su - M.u[0][1] * sv
// new sv = M.u[0][0] * sv - M.u[1][0] * su
// tmp1, tmp2: work buffers (each max(su.n, sv.n) + 2)
static void cofactor_apply_matrix1(
    SignedMpn& su, SignedMpn& sv,
    const mpn::HgcdMatrix1& M1,
    uint64_t* new_su_buf, uint64_t* new_sv_buf,
    uint64_t* tmp)
{
    size_t su_n, sv_n;
    int su_s, sv_s;

    // new su = u11 * su - u01 * sv
    smn_linearcomb_1(new_su_buf, &su_n, &su_s,
                     M1.u[1][1], su, M1.u[0][1], sv, tmp);

    // new sv = u00 * sv - u10 * su  (su still holds the old value)
    smn_linearcomb_1(new_sv_buf, &sv_n, &sv_s,
                     M1.u[0][0], sv, M1.u[1][0], su, tmp);

    // Copy the results back
    std::memcpy(su.d, new_su_buf, su_n * sizeof(uint64_t));
    su.n = su_n; su.s = su_s;
    std::memcpy(sv.d, new_sv_buf, sv_n * sizeof(uint64_t));
    sv.n = sv_n; sv.s = sv_s;
}

// Apply an HGCD multi-precision matrix to the cofactor pair
// new su = M.p[1][1] * su - M.p[0][1] * sv
// new sv = M.p[0][0] * sv - M.p[1][0] * su
static void cofactor_apply_hgcd_matrix(
    SignedMpn& su, SignedMpn& sv,
    const mpn::HgcdMatrix& M,
    uint64_t* new_su_buf, uint64_t* new_sv_buf,
    uint64_t* tmp1, uint64_t* tmp2, uint64_t* mul_scratch)
{
    size_t su_n, sv_n;
    int su_s, sv_s;

    smn_linearcomb_mp(new_su_buf, &su_n, &su_s,
                      M.p[1][1], M.n, su,
                      M.p[0][1], M.n, sv,
                      tmp1, tmp2, mul_scratch);

    smn_linearcomb_mp(new_sv_buf, &sv_n, &sv_s,
                      M.p[0][0], M.n, sv,
                      M.p[1][0], M.n, su,
                      tmp1, tmp2, mul_scratch);

    std::memcpy(su.d, new_su_buf, su_n * sizeof(uint64_t));
    su.n = su_n; su.s = su_s;
    std::memcpy(sv.d, new_sv_buf, sv_n * sizeof(uint64_t));
    sv.n = sv_n; sv.s = sv_s;
}

// Cofactor-tracking version of hgcd_reduce
// Reduce (up, vp) via HGCD while simultaneously updating (su, sv)
// Return value: new size (0 = no change)
static size_t hgcd_reduce_cofactor(
    uint64_t* up, uint64_t* vp, size_t n,
    SignedMpn& su, SignedMpn& sv,
    size_t cof_alloc)  // capacity of the cofactor buffers
{
    auto& arena = getThreadArena();
    size_t mark = arena.mark();

    size_t un = mpn::normalized_size(up, n);
    size_t vn = mpn::normalized_size(vp, n);
    if (un == 0 || vn == 0 || n < mpn::HGCD_THRESHOLD) {
        arena.rewind(mark);
        return n;
    }
    n = std::max(un, vn);

    // Pad to the same size
    if (un < n) std::memset(up + un, 0, (n - un) * sizeof(uint64_t));
    if (vn < n) std::memset(vp + vn, 0, (n - vn) * sizeof(uint64_t));

    // Ensure up >= vp
    bool swapped = false;
    if (mpn::cmp(up, n, vp, n) < 0) {
        for (size_t i = 0; i < n; i++) std::swap(up[i], vp[i]);
        swapped = true;
    }

    // Swap the cofactors as well
    if (swapped) {
        // Swap the data, size, and sign of su and sv
        std::swap(su.d, sv.d);
        std::swap(su.n, sv.n);
        std::swap(su.s, sv.s);
    }

    size_t p = n / 2;

    // Initialize the HGCD matrix
    mpn::HgcdMatrix M;
    mpn::hgcd_matrix_init(&M, 2 * (n - p) + 4);

    uint64_t* tp = arena.alloc_limbs(std::max(n - p, M.alloc) + 4);

    // Run HGCD on the top n-p limbs
    size_t nn = mpn::hgcd(up + p, vp + p, n - p, &M, tp);

    if (nn > 0) {
        // Apply the matrix to (up, vp)
        size_t mn = M.n;
        size_t prod_n = mn + p;
        size_t ms_sz = mpn::multiply_scratch_size(
            std::max(mn, p), std::min(mn, p));
        size_t adj_size = 2 * prod_n + ms_sz;
        uint64_t* adj_tp = arena.alloc_limbs(adj_size + 16);
        n = mpn::hgcd_matrix_adjust(&M, p + nn, up, vp, p, adj_tp);

        // Apply the matrix to the cofactors
        size_t cof_max = std::max(su.n, sv.n);
        size_t prod_cof = mn + cof_max + 4;
        size_t ms_cof = mpn::multiply_scratch_size(
            std::max(mn, cof_max), std::min({mn, cof_max, (size_t)1}));
        // Reserve enough for ms_cof
        ms_cof = std::max(ms_cof, mpn::multiply_scratch_size(mn + cof_max, 1));

        uint64_t* new_su = arena.alloc_limbs(prod_cof + 4);
        uint64_t* new_sv = arena.alloc_limbs(prod_cof + 4);
        uint64_t* ctmp1 = arena.alloc_limbs(prod_cof + 4);
        uint64_t* ctmp2 = arena.alloc_limbs(prod_cof + 4);
        uint64_t* cms = arena.alloc_limbs(ms_cof + 16);

        cofactor_apply_hgcd_matrix(su, sv, M, new_su, new_sv, ctmp1, ctmp2, cms);
    }

    n = std::max(mpn::normalized_size(up, n), mpn::normalized_size(vp, n));

    arena.rewind(mark);
    return n;
}

// subdiv step of gcdext (one subtraction + division, with cofactor update)
// up, vp: values (updated in-place)
// su, sv: cofactors (updated in-place)
// Return value: new size (0 = GCD found, stored in gp)
//
// Note: up, vp are passed by pointer value, so pointer swaps inside are not
// reflected to the caller. su, sv are passed by reference, so they are reflected.
// Consistency with the caller is maintained by swapping the physical buffer data.

// Signed addition: sv -= su  (in-place)
static void smn_subtract_inplace(SignedMpn& sv, const SignedMpn& su) {
    if (su.n == 0) return;
    if (sv.n == 0) {
        std::memcpy(sv.d, su.d, su.n * sizeof(uint64_t));
        sv.n = su.n;
        sv.s = -su.s;
        return;
    }
    if (sv.s != su.s) {
        // Different sign → add absolute values
        if (sv.n >= su.n) {
            uint64_t cy = mpn::add(sv.d, sv.d, sv.n, su.d, su.n);
            if (cy) sv.d[sv.n++] = cy;
        } else {
            uint64_t cy = mpn::add(sv.d, su.d, su.n, sv.d, sv.n);
            sv.n = su.n;
            if (cy) sv.d[sv.n++] = cy;
        }
    } else {
        // Same sign → subtract absolute values
        int c = (sv.n != su.n) ? ((sv.n > su.n) ? 1 : -1)
                                : mpn::cmp(sv.d, sv.n, su.d, su.n);
        if (c == 0) {
            sv.n = 0; sv.s = 0;
        } else if (c > 0) {
            mpn::sub(sv.d, sv.d, sv.n, su.d, su.n);
            sv.n = mpn::normalized_size(sv.d, sv.n);
        } else {
            mpn::sub(sv.d, su.d, su.n, sv.d, sv.n);
            sv.n = mpn::normalized_size(sv.d, su.n);
            sv.s = -sv.s;
        }
    }
}

// Physically swap the buffer data and the cofactors
static void phys_swap(uint64_t* up, uint64_t* vp, size_t n,
                      SignedMpn& su, SignedMpn& sv) {
    for (size_t i = 0; i < n; i++) std::swap(up[i], vp[i]);
    std::swap(su.d, sv.d); std::swap(su.n, sv.n); std::swap(su.s, sv.s);
}

static size_t gcdext_subdiv_step(
    uint64_t* up, uint64_t* vp, size_t n,
    SignedMpn& su, SignedMpn& sv,
    uint64_t* gp, size_t* gn,
    uint64_t* /*scratch*/, size_t /*cof_alloc*/)
{
    size_t an = mpn::normalized_size(up, n);
    size_t bn = mpn::normalized_size(vp, n);

    if (an == 0) {
        std::memcpy(gp, vp, bn * sizeof(uint64_t));
        *gn = bn;
        std::swap(su.d, sv.d); std::swap(su.n, sv.n); std::swap(su.s, sv.s);
        return 0;
    }
    if (bn == 0) {
        std::memcpy(gp, up, an * sizeof(uint64_t));
        *gn = an;
        return 0;
    }

    // Arrange up < vp (swap the physical buffers)
    if (an == bn) {
        int c = mpn::cmp(up, an, vp, bn);
        if (c == 0) {
            std::memcpy(gp, up, an * sizeof(uint64_t));
            *gn = an;
            return 0;
        }
        if (c > 0) {
            phys_swap(up, vp, n, su, sv);
            std::swap(an, bn);
        }
    } else if (an > bn) {
        phys_swap(up, vp, n, su, sv);
        std::swap(an, bn);
    }

    // vp > up: vp -= up (subtraction with q=1)
    mpn::sub(vp, vp, bn, up, an);
    bn = mpn::normalized_size(vp, bn);
    if (bn == 0) {
        std::memcpy(gp, up, an * sizeof(uint64_t));
        *gn = an;
        return 0;
    }

    // Cofactor update: sv -= su (corresponds to vp -= up)
    smn_subtract_inplace(sv, su);

    // Re-normalize sizes
    an = mpn::normalized_size(up, n);
    bn = mpn::normalized_size(vp, n);
    if (an == 0 || bn == 0) {
        if (bn == 0) {
            std::memcpy(gp, up, an * sizeof(uint64_t));
            *gn = an;
        } else {
            std::memcpy(gp, vp, bn * sizeof(uint64_t));
            *gn = bn;
            std::swap(su.d, sv.d); std::swap(su.n, sv.n); std::swap(su.s, sv.s);
        }
        return 0;
    }

    // Re-ensure up < vp (physical swap)
    if (an == bn) {
        int c = mpn::cmp(up, an, vp, bn);
        if (c == 0) {
            std::memcpy(gp, up, an * sizeof(uint64_t));
            *gn = an;
            return 0;
        }
        if (c > 0) {
            phys_swap(up, vp, n, su, sv);
            std::swap(an, bn);
        }
    } else if (an > bn) {
        phys_swap(up, vp, n, su, sv);
        std::swap(an, bn);
    }

    // vp = vp mod up (division)
    auto& arena = getThreadArena();
    size_t div_mark = arena.mark();

    size_t qn = bn - an + 1;
    uint64_t* qp = arena.alloc_limbs(qn + 4);
    uint64_t* dividend = arena.alloc_limbs(bn + 4);
    std::memcpy(dividend, vp, bn * sizeof(uint64_t));

    size_t ds = mpn::divide_scratch_size(bn, an);
    uint64_t* work = arena.alloc_limbs(ds + 4);

    mpn::divide(qp, vp, dividend, bn, up, an, work);

    // vp is the remainder (at most an limbs)
    if (an < n) std::memset(vp + an, 0, (n - an) * sizeof(uint64_t));
    size_t rem_n = mpn::normalized_size(vp, an);

    qn = mpn::normalized_size(qp, qn);

    if (rem_n == 0) {
        std::memcpy(gp, up, an * sizeof(uint64_t));
        *gn = an;
        arena.rewind(div_mark);
        return 0;
    }

    // Cofactor update: sv -= q * su
    if (qn > 0 && su.n > 0) {
        size_t prod_n = qn + su.n;
        uint64_t* prod = arena.alloc_limbs(prod_n + 4);

        if (qn == 1) {
            uint64_t cy = mpn::mul_1(prod, su.d, su.n, qp[0]);
            prod_n = su.n;
            if (cy) prod[prod_n++] = cy;
        } else {
            size_t ms = mpn::multiply_scratch_size(
                std::max(qn, su.n), std::min(qn, su.n));
            uint64_t* ms_buf = arena.alloc_limbs(ms + 4);
            if (qn >= su.n)
                mpn::multiply(prod, qp, qn, su.d, su.n, ms_buf);
            else
                mpn::multiply(prod, su.d, su.n, qp, qn, ms_buf);
            prod_n = mpn::normalized_size(prod, prod_n);
        }

        // sv -= su.s * prod
        if (sv.n == 0) {
            std::memcpy(sv.d, prod, prod_n * sizeof(uint64_t));
            sv.n = prod_n;
            sv.s = -su.s;
        } else if (sv.s != su.s) {
            // Different sign → add
            if (sv.n >= prod_n) {
                uint64_t cy = mpn::add(sv.d, sv.d, sv.n, prod, prod_n);
                if (cy) sv.d[sv.n++] = cy;
            } else {
                uint64_t cy = mpn::add(sv.d, prod, prod_n, sv.d, sv.n);
                sv.n = prod_n;
                if (cy) sv.d[sv.n++] = cy;
            }
        } else {
            // Same sign → subtract
            int c = (sv.n != prod_n) ? ((sv.n > prod_n) ? 1 : -1)
                                      : mpn::cmp(sv.d, sv.n, prod, prod_n);
            if (c == 0) {
                sv.n = 0; sv.s = 0;
            } else if (c > 0) {
                mpn::sub(sv.d, sv.d, sv.n, prod, prod_n);
                sv.n = mpn::normalized_size(sv.d, sv.n);
            } else {
                mpn::sub(sv.d, prod, prod_n, sv.d, sv.n);
                sv.n = mpn::normalized_size(sv.d, prod_n);
                sv.s = -sv.s;
            }
        }
    }

    arena.rewind(div_mark);
    return an;
}

} // anonymous namespace

// =========================================================================
// extendedGcd: HGCD + Lehmer fast version
// =========================================================================

Int IntGCD::extendedGcd(const Int& a, const Int& b, Int& x, Int& y) {
    // 1. Handle special states
    if (a.isNaN() || b.isNaN()) [[unlikely]] {
        x = Int::NaN(); y = Int::NaN(); return Int::NaN();
    }
    if (a.isInfinite() || b.isInfinite()) [[unlikely]] {
        x = Int::NaN(); y = Int::NaN(); return Int::NaN();
    }

    // 2. Handle zero cases
    if (b.isZero()) {
        if (a.isZero()) {
            x = Int::Zero(); y = Int::Zero(); return Int::Zero();
        }
        x = (a.getSign() >= 0) ? Int::One() : Int(-1);
        y = Int::Zero();
        return abs(a);
    }
    if (a.isZero()) {
        x = Int::Zero();
        y = (b.getSign() >= 0) ? Int::One() : Int(-1);
        return abs(b);
    }

    // 3. Remember the signs and compute with absolute values
    int a_sign = a.getSign();
    int b_sign = b.getSign();
    Int abs_a = abs(a);
    Int abs_b = abs(b);

    size_t an = abs_a.size();
    size_t bn = abs_b.size();

    // =========================================================
    // 4-5. mpn-level Lehmer + HGCD (unified across all sizes)
    //      Small sizes are handled by the n=1 fast loop
    // =========================================================
    ScratchScope scope;
    size_t max_n = std::max(an, bn);
    size_t buf_size = 2 * max_n + 32;

    auto& arena = getThreadArena();

    // Value buffers
    uint64_t* up = arena.alloc_limbs(buf_size);
    uint64_t* vp = arena.alloc_limbs(buf_size);
    uint64_t* gp = arena.alloc_limbs(buf_size);
    uint64_t* tp = arena.alloc_limbs(buf_size);

    std::memset(up, 0, buf_size * sizeof(uint64_t));
    std::memset(vp, 0, buf_size * sizeof(uint64_t));
    std::memset(gp, 0, buf_size * sizeof(uint64_t));

    for (size_t i = 0; i < an; i++) up[i] = abs_a.word(i);
    for (size_t i = 0; i < bn; i++) vp[i] = abs_b.word(i);

    // Arrange up >= vp (perform an initial division if unbalanced)
    size_t n;
    bool initial_swap = false;
    if (an > bn) {
        // Initial reduction via up mod vp
        size_t qn = an - bn + 1;
        uint64_t* qp_init = arena.alloc_limbs(qn + 4);
        uint64_t* tmp_init = arena.alloc_limbs(an + 4);
        std::memcpy(tmp_init, up, an * sizeof(uint64_t));
        size_t ds = mpn::divide_scratch_size(an, bn);
        uint64_t* ws = arena.alloc_limbs(ds + 4);
        mpn::divide(qp_init, up, tmp_init, an, vp, bn, ws);
        // up is the remainder
        bool zero = true;
        for (size_t i = 0; i < bn; i++) {
            if (up[i] != 0) { zero = false; break; }
        }
        if (zero) {
            // up was divisible by vp → gcd = vp
            Int g_result = Int::fromRawWords(
                std::span<const uint64_t>(vp, bn), +1);
            // |a| * x + |b| * y = gcd, x = 0, y = 1
            x = Int::Zero();
            y = Int::One();
            if (a_sign < 0) x = -x;
            if (b_sign < 0) y = -y;
            return g_result;
        }
        n = bn;
    } else if (an < bn) {
        // vp is larger → swap
        for (size_t i = 0; i < bn; i++) std::swap(up[i], vp[i]);
        initial_swap = true;

        size_t qn = bn - an + 1;
        uint64_t* qp_init = arena.alloc_limbs(qn + 4);
        uint64_t* tmp_init = arena.alloc_limbs(bn + 4);
        std::memcpy(tmp_init, up, bn * sizeof(uint64_t));
        size_t ds = mpn::divide_scratch_size(bn, an);
        uint64_t* ws = arena.alloc_limbs(ds + 4);
        mpn::divide(qp_init, up, tmp_init, bn, vp, an, ws);
        bool zero = true;
        for (size_t i = 0; i < an; i++) {
            if (up[i] != 0) { zero = false; break; }
        }
        if (zero) {
            Int g_result = Int::fromRawWords(
                std::span<const uint64_t>(vp, an), +1);
            x = Int::One();
            y = Int::Zero();
            if (initial_swap) std::swap(x, y);
            if (a_sign < 0) x = -x;
            if (b_sign < 0) y = -y;
            return g_result;
        }
        n = an;
    } else {
        n = an;
        // Ensure up >= vp
        if (mpn::cmp(up, n, vp, n) < 0) {
            for (size_t i = 0; i < n; i++) std::swap(up[i], vp[i]);
            initial_swap = true;
        }
    }

    // Allocate the cofactor buffers
    // The cofactors grow up to roughly the input size at most
    size_t cof_alloc = max_n + 16;
    uint64_t* su_buf = arena.alloc_limbs(cof_alloc);
    uint64_t* sv_buf = arena.alloc_limbs(cof_alloc);
    std::memset(su_buf, 0, cof_alloc * sizeof(uint64_t));
    std::memset(sv_buf, 0, cof_alloc * sizeof(uint64_t));

    SignedMpn su, sv;
    su.d = su_buf; su.d[0] = 1; su.n = 1; su.s = 1;  // su = 1 (cofactor of up)
    sv.d = sv_buf; sv.n = 0; sv.s = 0;                // sv = 0 (cofactor of vp)

    // Reflect the initial swap: in the initial_swap case, up is the original b and vp is the original a
    if (initial_swap) {
        // Since up = orig_b, vp = orig_a, su is the cofactor of orig_b
        // but the initial division has made up = orig_b mod orig_a
        // → su = 1 is the orig_b side of "up = orig_b * 1 + orig_a * (-q)"
        // We need to track the q of the initial division
        // → Simplification: in the initial_swap case, recompute without the swap
    }

    // ---- Lehmer + HGCD loop ----
    // Scratch for the Lehmer loop
    size_t lehmer_scratch = mpn::gcd_lehmer_scratch_size(n, n);
    uint64_t* subdiv_scratch = arena.alloc_limbs(lehmer_scratch + 16);

    size_t gn = 0;

    // Work buffer for the cofactor update (for 1-limb q × su/sv)
    uint64_t* cof_work = arena.alloc_limbs(cof_alloc + 4);

    while (n > 0) {
        // === n == 1: fast 1-limb loop that needs no arena ===
        if (n == 1) {
            uint64_t uu = up[0], vv = vp[0];
            if (uu == 0) {
                gp[0] = vv; gn = (vv > 0) ? 1 : 0;
                std::swap(su.d, sv.d); std::swap(su.n, sv.n); std::swap(su.s, sv.s);
                goto done;
            }
            if (vv == 0) {
                gp[0] = uu; gn = (uu > 0) ? 1 : 0;
                goto done;
            }

            // 1-limb Euclidean loop (with cofactor tracking)
            while (vv != 0) {
                // Ensure u >= v
                if (uu < vv) {
                    std::swap(uu, vv);
                    std::swap(su.d, sv.d); std::swap(su.n, sv.n); std::swap(su.s, sv.s);
                }

                uint64_t q = uu / vv;
                uint64_t r = uu - q * vv;
                uu = vv; vv = r;

                // Cofactor update: r = u - q*v → new_su = sv_old, new_sv = su_old - q*sv_old
                // Procedure: first swap (su = sv_old, sv = su_old), then sv -= q * su
                std::swap(su.d, sv.d); std::swap(su.n, sv.n); std::swap(su.s, sv.s);

                if (q > 0 && su.n > 0) {
                    uint64_t cy = mpn::mul_1(cof_work, su.d, su.n, q);
                    size_t pn = su.n;
                    if (cy) cof_work[pn++] = cy;

                    // sv -= su.s * cof_work[0..pn-1]
                    if (sv.n == 0) {
                        std::memcpy(sv.d, cof_work, pn * sizeof(uint64_t));
                        sv.n = pn; sv.s = -su.s;
                    } else if (sv.s != su.s) {
                        // Different sign → add
                        if (sv.n >= pn) {
                            uint64_t c = mpn::add(sv.d, sv.d, sv.n, cof_work, pn);
                            if (c) sv.d[sv.n++] = c;
                        } else {
                            uint64_t c = mpn::add(sv.d, cof_work, pn, sv.d, sv.n);
                            sv.n = pn;
                            if (c) sv.d[sv.n++] = c;
                        }
                    } else {
                        // Same sign → subtract
                        int c = (sv.n != pn) ? ((sv.n > pn) ? 1 : -1)
                                              : mpn::cmp(sv.d, sv.n, cof_work, pn);
                        if (c == 0) { sv.n = 0; sv.s = 0; }
                        else if (c > 0) {
                            mpn::sub(sv.d, sv.d, sv.n, cof_work, pn);
                            sv.n = mpn::normalized_size(sv.d, sv.n);
                        } else {
                            mpn::sub(sv.d, cof_work, pn, sv.d, sv.n);
                            sv.n = mpn::normalized_size(sv.d, pn);
                            sv.s = -sv.s;
                        }
                    }
                }
            }

            gp[0] = uu; gn = (uu > 0) ? 1 : 0;
            goto done;
        }

        // === n >= 2: fast reduction via HGCD
        if (n >= mpn::HGCD_THRESHOLD) {
            size_t new_n = hgcd_reduce_cofactor(up, vp, n, su, sv, cof_alloc);
            if (new_n < n) {
                n = new_n;
                size_t un2 = mpn::normalized_size(up, n);
                size_t vn2 = mpn::normalized_size(vp, n);
                if (un2 == 0) {
                    std::memcpy(gp, vp, vn2 * sizeof(uint64_t));
                    gn = vn2;
                    std::swap(su.d, sv.d); std::swap(su.n, sv.n); std::swap(su.s, sv.s);
                    goto done;
                }
                if (vn2 == 0) {
                    std::memcpy(gp, up, un2 * sizeof(uint64_t));
                    gn = un2;
                    goto done;
                }
                continue;
            }
        }

        // Lehmer step via hgcd2
        mpn::HgcdMatrix1 M1;
        uint64_t mask = up[n - 1] | vp[n - 1];
        if (mask == 0) { n--; continue; }

        uint64_t uh, ul, vh, vl;
        if (mask >> 63) {
            uh = up[n-1]; ul = up[n-2];
            vh = vp[n-1]; vl = vp[n-2];
        } else {
            int shift = std::countl_zero(mask);
            if (n >= 3) {
                uh = (up[n-1] << shift) | (up[n-2] >> (64 - shift));
                ul = (up[n-2] << shift) | (up[n-3] >> (64 - shift));
                vh = (vp[n-1] << shift) | (vp[n-2] >> (64 - shift));
                vl = (vp[n-2] << shift) | (vp[n-3] >> (64 - shift));
            } else {
                uh = (up[n-1] << shift) | (up[n-2] >> (64 - shift));
                ul = up[n-2] << shift;
                vh = (vp[n-1] << shift) | (vp[n-2] >> (64 - shift));
                vl = vp[n-2] << shift;
            }
        }

        if (mpn::hgcd2(uh, ul, vh, vl, &M1)) {
            // Apply the matrix to (up, vp)
            n = mpn::hgcd_mul_matrix1_vector(&M1, tp, up, vp, n);
            std::swap(up, tp);

            // Apply the same matrix to the cofactors
            size_t cm = std::max(su.n, sv.n) + 2;
            size_t ctmp_mark = arena.mark();
            uint64_t* new_su = arena.alloc_limbs(cm + 4);
            uint64_t* new_sv = arena.alloc_limbs(cm + 4);
            uint64_t* ctmp = arena.alloc_limbs(cm + 4);
            cofactor_apply_matrix1(su, sv, M1, new_su, new_sv, ctmp);
            arena.rewind(ctmp_mark);

            if (n == 0) { gp[0] = 0; gn = 0; goto done; }
        } else {
            // subdiv step (subtraction + division + cofactor update)
            n = gcdext_subdiv_step(up, vp, n, su, sv, gp, &gn,
                                   subdiv_scratch, cof_alloc);
            if (n == 0) goto done;
        }
    }

    // We only reach here when n == 0 (normally we exit via goto done)
    gp[0] = 0; gn = 0;

done:
    // gp[0..gn-1] holds the GCD, su holds the cofactor of up
    if (gn == 0) {
        x = Int::Zero(); y = Int::Zero(); return Int::Zero();
    }

    Int g_result = Int::fromRawWords(std::span<const uint64_t>(gp, gn), +1);

    // su → x_cof (cofactor of up)
    Int x_cof;
    if (su.n == 0) {
        x_cof = Int::Zero();
    } else {
        x_cof = Int::fromRawWords(std::span<const uint64_t>(su.d, su.n), su.s);
    }

    if (initial_swap) {
        y = x_cof;
        x = (g_result - abs_b * y) / abs_a;
    } else {
        x = x_cof;
        y = (g_result - abs_a * x) / abs_b;
    }

    if (a_sign < 0) x = -x;
    if (b_sign < 0) y = -y;
    return g_result;
}

} // namespace sangi
