namespace PerfectNumbers.Core;

internal struct ExponentRemainderStepper(in MontgomeryDivisorData divisor)
{
    private readonly MontgomeryDivisorData _divisor = divisor;
    private readonly ulong _modulus = divisor.Modulus;
    private readonly ulong _nPrime = divisor.NPrime;
    private readonly ulong _montgomeryOne = divisor.MontgomeryOne;
    private ulong _previousExponent = 0UL;
    private ulong _currentMontgomery = divisor.MontgomeryOne;
    private bool _hasState = false;

	public readonly bool IsValidModulus => _modulus > 1UL && (_modulus & 1UL) != 0UL;

    public readonly bool HasState => _hasState;

    public readonly ulong PreviousExponent => _previousExponent;

    public void Reset()
    {
        _previousExponent = 0UL;
        _currentMontgomery = _montgomeryOne;
        _hasState = false;
    }

    public ulong ComputeNext(ulong exponent)
    {
        if (!IsValidModulus)
        {
            _hasState = false;
            return 0UL;
        }

        if (!_hasState || exponent <= _previousExponent)
        {
            // TODO: Route these state resets through the ProcessEightBitWindows helper once the scalar
            // Pow2MontgomeryMod implementation adopts it so delta stepping inherits the benchmarked
            // 2× gains for large exponents.
            _currentMontgomery = exponent.Pow2MontgomeryModWindowedCpu(_divisor, keepMontgomery: true);
            _previousExponent = exponent;
            _hasState = true;
            return ReduceCurrent();
        }

        ulong delta = exponent - _previousExponent;
        // TODO: Once divisor cycle lengths are mandatory, pull the delta multiplier from the
        // single-block divisor-cycle snapshot so we can skip the powmod entirely and reuse the
        // cached Montgomery residue ladder highlighted in MersenneDivisorCycleLengthGpuBenchmarks.
        // TODO: Replace this per-delta powmod with the upcoming windowed ladder so incremental
        // updates stop paying the single-bit cost highlighted in GpuPow2ModBenchmarks.
        ulong multiplier = delta.Pow2MontgomeryModWindowedCpu(_divisor, keepMontgomery: true);
        _currentMontgomery = _currentMontgomery.MontgomeryMultiply(multiplier, _modulus, _nPrime);
        _previousExponent = exponent;
        return ReduceCurrent();
    }

    public bool ComputeNextIsUnity(ulong exponent)
    {
        if (!IsValidModulus)
        {
            _hasState = false;
            return false;
        }

        if (!_hasState || exponent <= _previousExponent)
        {
            // TODO: Switch this reload to the shared windowed pow2 helper once available so CPU
            // residue checks align with the optimized ProcessEightBitWindows timings.
            _currentMontgomery = exponent.Pow2MontgomeryModWindowedCpu(_divisor, keepMontgomery: true);
            _previousExponent = exponent;
            _hasState = true;
            return _currentMontgomery == _montgomeryOne;
        }

        ulong delta = exponent - _previousExponent;
        // TODO: Reuse the divisor-cycle derived Montgomery delta once the cache exposes single-cycle
        // lookups so this branch also avoids recomputing powmods when the snapshot lacks the divisor.
        // TODO: Use the windowed delta pow2 helper here as well to avoid the single-bit ladder that
        // currently lags behind the benchmarked implementation for huge divisor cycles.
        ulong multiplier = delta.Pow2MontgomeryModWindowedCpu(_divisor, keepMontgomery: true);
        _currentMontgomery = _currentMontgomery.MontgomeryMultiply(multiplier, _modulus, _nPrime);
        _previousExponent = exponent;
        return _currentMontgomery == _montgomeryOne;
    }

    public bool TryInitializeFromMontgomeryResult(ulong exponent, ulong montgomeryResult, out bool isUnity)
    {
        if (!IsValidModulus)
        {
            _hasState = false;
            isUnity = false;
            return false;
        }

        _currentMontgomery = montgomeryResult;
        _previousExponent = exponent;
        _hasState = true;
        isUnity = _currentMontgomery == _montgomeryOne;
        return true;
    }

    public bool TryAdvanceWithMontgomeryDelta(ulong exponent, ulong montgomeryDelta, out bool isUnity)
    {
        if (!IsValidModulus)
        {
            _hasState = false;
            isUnity = false;
            return false;
        }

        if (!_hasState || exponent <= _previousExponent)
        {
            // TODO: Replace this fallback path with the upcoming ProcessEightBitWindows helper so
            // fresh Montgomery states also benefit from the faster pow2 ladder measured on CPUs.
            _currentMontgomery = exponent.Pow2MontgomeryModWindowedCpu(_divisor, keepMontgomery: true);
            _previousExponent = exponent;
            _hasState = true;
            isUnity = _currentMontgomery == _montgomeryOne;
            return true;
        }

        // TODO: Once the divisor-cycle cache exposes a direct Montgomery delta, multiply it here instead
        // of relying on the caller-provided delta so incremental scans remain in sync with the snapshot
        // without computing additional cycles or mutating cache state.
        _currentMontgomery = _currentMontgomery.MontgomeryMultiply(montgomeryDelta, _modulus, _nPrime);
        _previousExponent = exponent;
        isUnity = _currentMontgomery == _montgomeryOne;
        return true;
    }

    private ulong ReduceCurrent() => _currentMontgomery.MontgomeryMultiply(1UL, _modulus, _nPrime);
}