using System.Buffers;
using System.Runtime.InteropServices;
using ILGPU.Runtime;

namespace PerfectNumbers.Core.Gpu;

/// <summary>
/// Uses GPU kernels to incrementally scan Mersenne candidates for prime exponents p >= 31.
/// </summary>
public class MersenneNumberIncrementalGpuTester(GpuKernelType kernelType, bool useGpuOrder)
{
    private readonly GpuKernelType _kernelType = kernelType;
    private readonly bool _useGpuOrder = useGpuOrder;

	private static readonly Accelerator[] _accelerators = AcceleratorPool.Shared.Accelerators;

    public void Scan(ulong exponent, UInt128 twoP, LastDigit lastDigit, UInt128 maxK, ref bool isPrime)
    {
		// TODO: Replace LastDigit lastDigit with bool lastIsSeven
        throw new NotImplementedException($"GPU incremental scanning requires the device cycle heuristics implementation (kernel: {_kernelType}, GPU order: {_useGpuOrder}).");

        // GpuPrimeWorkLimiter.Acquire();
        var acceleratorIndex = AcceleratorPool.Shared.Rent();
		var accelerator = _accelerators[acceleratorIndex];
        var stream = accelerator.CreateStream();
        int batchSize = GpuConstants.ScanBatchSize; // large batch improves GPU occupancy
        UInt128 kStart = 1UL;
        ulong divMul = (ulong)((((UInt128)1 << 64) - UInt128.One) / exponent) + 1UL;
        byte last = lastDigit == LastDigit.Seven ? (byte)1 : (byte)0; // ILGPU kernels do not support bool parameters

		KernelContainer kernels = GpuKernelPool.GetOrAddKernels(acceleratorIndex, stream, KernelType.Pow2ModKernelScan | KernelType.IncrementalKernelScan);
		var pow2Kernel = kernels.Pow2Mod!;
        var incKernel = kernels.Incremental!;
        ulong step10 = ((ulong)(exponent.Mod10() << 1)).Mod10();
        ulong step8 = ((exponent & 7UL) << 1) & 7UL;
        ulong step3 = ((exponent % 3UL) << 1) % 3UL;
        ulong step5 = ((exponent % 5UL) << 1) % 5UL;
        GpuUInt128 twoPGpu = (GpuUInt128)twoP;
        var smallCyclesView = GpuKernelPool.GetSmallCyclesOnDevice(kernels);
        ResiduePrimeViews primeViews = default;
        if (_kernelType == GpuKernelType.Pow2Mod)
        {
            primeViews = GpuKernelPool.GetSmallPrimesOnDevice(kernels);
        }

        var orderBuffer = accelerator.Allocate1D<ulong>(batchSize);
		// Avoid giant stack allocations that can trigger StackOverflow when batchSize is large.
		// Rent a reusable array from the shared pool instead.
		ArrayPool<ulong> ulongPool = ThreadStaticPools.UlongPool;
		ulong[] orderArray = ulongPool.Rent(batchSize);
        UInt128 remaining;
        int currentSize;
        int i;
        UInt128 q = UInt128.Zero;
        while (kStart <= maxK && Volatile.Read(ref isPrime))
        {
            remaining = maxK - kStart + UInt128.One;
            currentSize = remaining > (UInt128)batchSize ? batchSize : (int)remaining;
            Span<ulong> orders = orderArray.AsSpan(0, currentSize);
            // Precompute residue automaton bases for this batch
            UInt128 q0 = twoP * kStart + UInt128.One;
            if (_kernelType == GpuKernelType.Pow2Mod)
            {
                q0.Mod10_8_5_3(out ulong q0m10, out ulong q0m8, out ulong q0m5, out ulong q0m3);
                var kernelArgs = new ResidueAutomatonArgs(q0m10, step10, q0m8, step8, q0m3, step3, q0m5, step5);
                pow2Kernel(
                    stream,
                    currentSize,
                    exponent,
                    twoPGpu,
                    (GpuUInt128)kStart,
                    last,
                    divMul,
                    kernelArgs,
                    orderBuffer.View,
                    smallCyclesView,
                    primeViews.LastOne,
                    primeViews.LastSeven,
                    primeViews.LastOnePow2,
                    primeViews.LastSevenPow2);
            }
            else
            {
                q0.Mod10_8_5_3(out ulong q0m10, out ulong q0m8, out ulong q0m5, out ulong q0m3);
                incKernel(
                    stream,
                    currentSize,
                    exponent,
                    twoPGpu,
                    (GpuUInt128)kStart,
                    last,
                    divMul,
                    q0m10,
                    q0m8,
                    q0m3,
                    q0m5,
                    orderBuffer.View,
                    smallCyclesView);
            }

            orderBuffer.View.CopyToCPU(stream, ref MemoryMarshal.GetReference(orders), currentSize);
            stream.Synchronize();
            if (_kernelType == GpuKernelType.Pow2Mod)
            {
                for (i = 0; i < currentSize; i++)
                {
                    if (orders[i] != 0UL)
                    {
                        Volatile.Write(ref isPrime, false);
                        break;
                    }
                }
            }
            else
            {
                q = twoP * kStart + UInt128.One;
                for (i = 0; i < currentSize && Volatile.Read(ref isPrime); i++, q += twoP)
                {
                    if (orders[i] != 0UL && q.IsPrimeCandidate())
                    {
                        Volatile.Write(ref isPrime, false);
                    }
                }
            }

            kStart += (UInt128)currentSize;
        }

		stream.Dispose();
        ulongPool.Return(orderArray);
		orderBuffer.Dispose();
        // GpuPrimeWorkLimiter.Release();
    }
}