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

// IntRandom.hpp — Multi-precision random integer generation
//
// GMP-compatible API:
//   randomBits(n)       : uniform random in [0, 2^n) (mpz_urandomb)
//   randomBelow(upper)  : uniform random in [0, upper) (mpz_urandomm)

#ifndef SANGI_INT_RANDOM_HPP
#define SANGI_INT_RANDOM_HPP

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

namespace sangi {

/**
 * @brief Multi-precision random-integer generator
 *
 * Holds an internal mt19937_64 engine and generates uniform random Int of arbitrary bit width.
 * The default constructor seeds from std::random_device.
 */
class IntRandom {
public:
    /// Default constructor (seeded from random_device)
    IntRandom() {
        std::random_device rd;
        std::seed_seq seq{rd(), rd(), rd(), rd(), rd(), rd(), rd(), rd()};
        m_engine.seed(seq);
    }

    /// Explicit seed
    explicit IntRandom(uint64_t seed) : m_engine(seed) {}

    /// Generate a uniform random number in [0, 2^bits) (equivalent to GMP mpz_urandomb)
    /// bits == 0 -> return 0
    Int randomBits(size_t bits) {
        if (bits == 0) return Int(0);

        size_t nwords = (bits + 63) / 64;
        Int result;
        result.m_words.resize_uninitialized(nwords);

        for (size_t i = 0; i < nwords; ++i) {
            result.m_words[i] = m_engine();
        }

        // Mask the top word: clear high bits if bits is not a multiple of 64
        size_t topBits = bits % 64;
        if (topBits != 0) {
            result.m_words[nwords - 1] &= (uint64_t(1) << topBits) - 1;
        }

        // normalize: remove leading zero words
        while (!result.m_words.empty() && result.m_words.back() == 0) {
            result.m_words.pop_back();
        }

        result.m_sign = result.m_words.empty() ? 0 : 1;
        result.m_state = NumericState::Normal;
        return result;
    }

    /// Generate a uniform random number in [0, upper) (equivalent to GMP mpz_urandomm)
    /// Rejection sampling: generate a random number of upper.bitLength() bits and
    /// retry if it is >= upper
    /// upper <= 0 -> return 0
    Int randomBelow(const Int& upper) {
        if (upper.getSign() <= 0) return Int(0);

        size_t bits = upper.bitLength();

        // rejection sampling
        for (;;) {
            Int r = randomBits(bits);
            if (r < upper) return r;
        }
    }

    /// Generate a uniform random number in [lo, hi] (inclusive on both ends)
    Int randomRange(const Int& lo, const Int& hi) {
        if (lo > hi) return Int(0);
        if (lo == hi) return lo;
        Int range = hi - lo + 1;
        return lo + randomBelow(range);
    }

    /// Biased random for testing (equivalent to GMP mpz_rrandomb)
    /// Generate a positive integer of bits bits with long bit-runs (consecutive 0/1).
    /// MSB is always 1. Useful for edge-case testing of multiplication/division.
    /// bits == 0 -> return 0
    Int rrandomb(size_t bits) {
        if (bits == 0) return Int(0);

        size_t nwords = (bits + 63) / 64;
        Int result;
        result.m_words.resize_uninitialized(nwords);
        for (size_t i = 0; i < nwords; ++i)
            result.m_words[i] = 0;

        // Fill with random-length bit-runs (GMP style)
        // Toggle the current bit value (0 or 1)
        bool current_bit = true;  // Start from MSB side with 1
        size_t pos = bits;        // Remaining number of bits

        while (pos > 0) {
            // Determine run length geometrically (mostly short, occasionally long)
            // GMP: when gmp_urandomb_ui(state, 0) is used, BITS_PER_MP_LIMB is the upper bound
            size_t max_run = std::min(pos, static_cast<size_t>(64));
            uint64_t r = m_engine();
            // Take (count of consecutive 0s from the low bit + 1) as the run length (approximating a geometric distribution)
            size_t run_len;
            if (r == 0) {
                run_len = max_run;
            } else {
                run_len = static_cast<size_t>(std::countr_zero(r)) + 1;
                if (run_len > max_run) run_len = max_run;
            }

            // Fill bits [pos - run_len, pos) with current_bit
            if (current_bit) {
                for (size_t i = 0; i < run_len; ++i) {
                    size_t bit_pos = pos - 1 - i;
                    result.m_words[bit_pos / 64] |= uint64_t(1) << (bit_pos % 64);
                }
            }
            // If current_bit == false, the bits are already 0, so skip

            pos -= run_len;
            current_bit = !current_bit;
        }

        // Ensure the MSB is 1 (as a bit-bit number)
        result.m_words[(bits - 1) / 64] |= uint64_t(1) << ((bits - 1) % 64);

        // Clear surplus bits in the top word
        size_t topBits = bits % 64;
        if (topBits != 0) {
            result.m_words[nwords - 1] &= (uint64_t(1) << topBits) - 1;
        }

        // normalize
        while (!result.m_words.empty() && result.m_words.back() == 0)
            result.m_words.pop_back();

        result.m_sign = result.m_words.empty() ? 0 : 1;
        result.m_state = NumericState::Normal;
        return result;
    }

    /// Access to the engine (for external customization)
    std::mt19937_64& engine() { return m_engine; }

private:
    std::mt19937_64 m_engine;
};

/// Free functions using the thread_local global engine

/// Uniform random in [0, 2^bits)
inline Int randomBits(size_t bits) {
    thread_local IntRandom g_rng;
    return g_rng.randomBits(bits);
}

/// Uniform random in [0, upper)
inline Int randomBelow(const Int& upper) {
    thread_local IntRandom g_rng;
    return g_rng.randomBelow(upper);
}

/// Uniform random in [lo, hi]
inline Int randomRange(const Int& lo, const Int& hi) {
    thread_local IntRandom g_rng;
    return g_rng.randomRange(lo, hi);
}

/// Biased random for testing (long bit-runs)
inline Int rrandomb(size_t bits) {
    thread_local IntRandom g_rng;
    return g_rng.rrandomb(bits);
}

} // namespace sangi

#endif // SANGI_INT_RANDOM_HPP