using System;
using System.Buffers;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using System.Threading;
using ILGPU;
using ILGPU.Runtime;
using PerfectNumbers.Core.Gpu;

namespace PerfectNumbers.Core;

public sealed class DivisorCycleCache
{
    private const byte ByteZero = 0;
    private const byte ByteOne = 1;
    private const int StackBufferThreshold = 128;

    private readonly ConcurrentDictionary<Accelerator, Action<Index1D, int, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<byte, Stride1D.Dense>>> _gpuKernelCache = new(AcceleratorReferenceComparer.Instance);
    private volatile ulong[] _snapshot;
    private volatile bool _useGpuGeneration = true;
    private readonly int _divisorCyclesBatchSize;
    private static int _sharedDivisorCyclesBatchSize = GpuConstants.GpuCycleStepsPerInvocation;

    public static void SetDivisorCyclesBatchSize(int divisorCyclesBatchSize)
    {
        _sharedDivisorCyclesBatchSize = Math.Max(1, divisorCyclesBatchSize);
    }

    public static DivisorCycleCache Shared { get; } = new DivisorCycleCache(_sharedDivisorCyclesBatchSize);

    public int PreferredBatchSize => _divisorCyclesBatchSize;

    private DivisorCycleCache(int divisorCyclesBatchSize)
    {
        _divisorCyclesBatchSize = Math.Max(1, divisorCyclesBatchSize);
        _snapshot = MersenneDivisorCycles.Shared.ExportSmallCyclesSnapshot();
    }

    public void ConfigureGeneratorDevice(bool useGpu)
    {
        _useGpuGeneration = useGpu;
    }

    public void RefreshSnapshot()
    {
        _snapshot = MersenneDivisorCycles.Shared.ExportSmallCyclesSnapshot();
    }

    public ulong GetCycleLength(ulong divisor)
    {
        Span<ulong> result = stackalloc ulong[1];
        ReadOnlySpan<ulong> singleDivisor = stackalloc ulong[1] { divisor };
        GetCycleLengths(singleDivisor, result);
        return result[0];
    }

    public void GetCycleLengths(ReadOnlySpan<ulong> divisors, Span<ulong> cycles)
    {
        ulong[] snapshot = _snapshot;
        int length = divisors.Length;
        int[]? rentedMissing = null;
        Span<int> missingBuffer = length <= StackBufferThreshold
            ? stackalloc int[length]
            : length < PerfectNumberConstants.PooledArrayThreshold
                ? new Span<int>(new int[length])
                : new Span<int>(rentedMissing = ThreadLocalArrayPool<int>.Shared.Rent(length), 0, length);

        int missingCount = 0;

        for (int i = 0; i < length; i++)
        {
            ulong divisor = divisors[i];
            if (divisor <= 1UL)
            {
                throw new InvalidDataException("Divisor must be > 1");
            }

            if (divisor < (ulong)snapshot.Length)
            {
                ulong cached = snapshot[divisor];
                if (cached == 0UL)
                {
                    throw new InvalidDataException($"Divisor cycle is missing for {divisor}");
                }

                cycles[i] = cached;
            }
            else
            {
                missingBuffer[missingCount++] = i;
            }
        }

        if (missingCount != 0)
        {
            ReadOnlySpan<int> missingIndices = missingBuffer[..missingCount];
            if (_useGpuGeneration)
            {
                ComputeCyclesGpu(divisors, cycles, missingIndices);
            }
            else
            {
                ComputeCyclesCpu(divisors, cycles, missingIndices);
            }
        }

        if (rentedMissing is not null)
        {
            ThreadLocalArrayPool<int>.Shared.Return(rentedMissing, clearArray: false);
        }
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private static void ComputeCyclesCpu(ReadOnlySpan<ulong> divisors, Span<ulong> cycles, ReadOnlySpan<int> indices)
    {
        for (int i = 0; i < indices.Length; i++)
        {
            int targetIndex = indices[i];
            ulong divisor = divisors[targetIndex];
            MontgomeryDivisorData divisorData = MontgomeryDivisorData.FromModulus(divisor);
            cycles[targetIndex] = MersenneDivisorCycles.CalculateCycleLength(divisor, divisorData);
        }
    }

    private void ComputeCyclesGpu(ReadOnlySpan<ulong> divisors, Span<ulong> cycles, ReadOnlySpan<int> indices)
    {
        int count = indices.Length;
        ulong[]? rentedDivisors = null;
        Span<ulong> workDivisors = count <= StackBufferThreshold
            ? stackalloc ulong[count]
            : count < PerfectNumberConstants.PooledArrayThreshold
                ? new Span<ulong>(new ulong[count])
                : new Span<ulong>(rentedDivisors = ThreadLocalArrayPool<ulong>.Shared.Rent(count), 0, count);

        ulong[]? rentedResults = null;
        Span<ulong> workResults = count <= StackBufferThreshold
            ? stackalloc ulong[count]
            : count < PerfectNumberConstants.PooledArrayThreshold
                ? new Span<ulong>(new ulong[count])
                : new Span<ulong>(rentedResults = ThreadLocalArrayPool<ulong>.Shared.Rent(count), 0, count);

        for (int i = 0; i < count; i++)
        {
            workDivisors[i] = divisors[indices[i]];
        }

        ComputeCyclesGpuCore(workDivisors, workResults);

        for (int i = 0; i < count; i++)
        {
            cycles[indices[i]] = workResults[i];
        }

        if (rentedDivisors is not null)
        {
            ThreadLocalArrayPool<ulong>.Shared.Return(rentedDivisors, clearArray: false);
        }

        if (rentedResults is not null)
        {
            ThreadLocalArrayPool<ulong>.Shared.Return(rentedResults, clearArray: false);
        }
    }

    private void ComputeCyclesGpuCore(ReadOnlySpan<ulong> divisors, Span<ulong> destination)
    {
        int length = divisors.Length;
        var gpuLease = GpuKernelPool.GetKernel(useGpuOrder: true);
        var execution = gpuLease.EnterExecutionScope();

        Accelerator accelerator = gpuLease.Accelerator;
        var kernel = _gpuKernelCache.GetOrAdd(accelerator, LoadKernel);

        using MemoryBuffer1D<ulong, Stride1D.Dense> divisorBuffer = accelerator.Allocate1D<ulong>(length);
        using MemoryBuffer1D<ulong, Stride1D.Dense> powBuffer = accelerator.Allocate1D<ulong>(length);
        using MemoryBuffer1D<ulong, Stride1D.Dense> orderBuffer = accelerator.Allocate1D<ulong>(length);
        using MemoryBuffer1D<ulong, Stride1D.Dense> resultBuffer = accelerator.Allocate1D<ulong>(length);
        using MemoryBuffer1D<byte, Stride1D.Dense> statusBuffer = accelerator.Allocate1D<byte>(length);

        ulong[]? rentedPow = null;
        Span<ulong> powSpan = length <= StackBufferThreshold
            ? stackalloc ulong[length]
            : length < PerfectNumberConstants.PooledArrayThreshold
                ? new Span<ulong>(new ulong[length])
                : new Span<ulong>(rentedPow = ThreadLocalArrayPool<ulong>.Shared.Rent(length), 0, length);

        ulong[]? rentedOrder = null;
        Span<ulong> orderSpan = length <= StackBufferThreshold
            ? stackalloc ulong[length]
            : length < PerfectNumberConstants.PooledArrayThreshold
                ? new Span<ulong>(new ulong[length])
                : new Span<ulong>(rentedOrder = ThreadLocalArrayPool<ulong>.Shared.Rent(length), 0, length);

        ulong[]? rentedResult = null;
        Span<ulong> resultSpan = length <= StackBufferThreshold
            ? stackalloc ulong[length]
            : length < PerfectNumberConstants.PooledArrayThreshold
                ? new Span<ulong>(new ulong[length])
                : new Span<ulong>(rentedResult = ThreadLocalArrayPool<ulong>.Shared.Rent(length), 0, length);

        byte[]? rentedStatus = null;
        Span<byte> statusSpan = length <= StackBufferThreshold
            ? stackalloc byte[length]
            : length < PerfectNumberConstants.PooledArrayThreshold
                ? new Span<byte>(new byte[length])
                : new Span<byte>(rentedStatus = ThreadLocalArrayPool<byte>.Shared.Rent(length), 0, length);

        ref ulong divisorRef = ref MemoryMarshal.GetReference(divisors);
        divisorBuffer.View.CopyFromCPU(ref divisorRef, length);

        for (int i = 0; i < length; i++)
        {
            ulong divisor = divisors[i];
            ulong initialPow = 2UL;
            if (initialPow >= divisor)
            {
                initialPow -= divisor;
            }

            powSpan[i] = initialPow;
            orderSpan[i] = 1UL;
            resultSpan[i] = 0UL;
            statusSpan[i] = ByteZero;
        }

        powBuffer.View.CopyFromCPU(ref MemoryMarshal.GetReference(powSpan), length);
        orderBuffer.View.CopyFromCPU(ref MemoryMarshal.GetReference(orderSpan), length);
        resultBuffer.View.CopyFromCPU(ref MemoryMarshal.GetReference(resultSpan), length);
        statusBuffer.View.CopyFromCPU(ref MemoryMarshal.GetReference(statusSpan), length);

        int pending;
        do
        {
            kernel(length, _divisorCyclesBatchSize, divisorBuffer.View, powBuffer.View, orderBuffer.View, resultBuffer.View, statusBuffer.View);
            accelerator.Synchronize();

            statusBuffer.View.CopyToCPU(ref MemoryMarshal.GetReference(statusSpan), length);
            pending = 0;
            for (int i = 0; i < length; i++)
            {
                if (statusSpan[i] == ByteZero)
                {
                    pending++;
                }
            }
        }
        while (pending > 0);

        resultBuffer.View.CopyToCPU(ref MemoryMarshal.GetReference(resultSpan), length);
        resultSpan.CopyTo(destination);
        if (rentedPow is not null)
        {
            ThreadLocalArrayPool<ulong>.Shared.Return(rentedPow, clearArray: false);
        }

        if (rentedOrder is not null)
        {
            ThreadLocalArrayPool<ulong>.Shared.Return(rentedOrder, clearArray: false);
        }

        if (rentedResult is not null)
        {
            ThreadLocalArrayPool<ulong>.Shared.Return(rentedResult, clearArray: false);
        }

        if (rentedStatus is not null)
        {
            ThreadLocalArrayPool<byte>.Shared.Return(rentedStatus, clearArray: false);
        }

        execution.Dispose();
        gpuLease.Dispose();
    }

    private static Action<Index1D, int, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<byte, Stride1D.Dense>> LoadKernel(Accelerator accelerator)
    {
        return accelerator.LoadAutoGroupedStreamKernel<Index1D, int, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<byte, Stride1D.Dense>>(GpuAdvanceDivisorCyclesKernel);
    }

    private static void GpuAdvanceDivisorCyclesKernel(
        Index1D index,
        int steps,
        ArrayView1D<ulong, Stride1D.Dense> divisors,
        ArrayView1D<ulong, Stride1D.Dense> pow,
        ArrayView1D<ulong, Stride1D.Dense> order,
        ArrayView1D<ulong, Stride1D.Dense> cycles,
        ArrayView1D<byte, Stride1D.Dense> status)
    {
        if (status[index] != ByteZero)
        {
            return;
        }

        ulong divisor = divisors[index];

        if ((divisor & (divisor - 1UL)) == 0UL)
        {
            cycles[index] = 1UL;
            status[index] = ByteOne;
            pow[index] = 1UL;
            order[index] = 1UL;
            return;
        }

        ulong currentPow = pow[index];
        ulong currentOrder = order[index];

        if (divisor <= 3UL)
        {
            while (currentPow != 1UL)
            {
                currentPow += currentPow;
                if (currentPow >= divisor)
                {
                    currentPow -= divisor;
                }

                currentOrder++;
            }

            cycles[index] = currentOrder;
            status[index] = ByteOne;
            pow[index] = currentPow;
            order[index] = currentOrder;
            return;
        }

        do
        {
            currentPow += currentPow;
            if (currentPow >= divisor)
            {
                currentPow -= divisor;
            }

            currentOrder++;

            if (currentPow == 1UL)
            {
                cycles[index] = currentOrder;
                status[index] = ByteOne;
                pow[index] = currentPow;
                order[index] = currentOrder;
                return;
            }
        }
        while (--steps != 0);

        pow[index] = currentPow;
        order[index] = currentOrder;
    }

    private sealed class AcceleratorReferenceComparer : IEqualityComparer<Accelerator>
    {
        internal static AcceleratorReferenceComparer Instance { get; } = new();

        public bool Equals(Accelerator? x, Accelerator? y) => ReferenceEquals(x, y);

        public int GetHashCode(Accelerator obj) => RuntimeHelpers.GetHashCode(obj);
    }
}