XFESystem

using System.Buffers.Binary;

namespace XFESystem.Gpu.Nvidia;

/// <summary>
/// Native NVIDIA method/GPFIFO encoding, derived from the MIT-licensed public
/// headers pinned in Tools/Gpu/NvidiaSources.lock.json. This emits GPU commands;
/// it does not emulate, execute or claim CUDA compatibility.
/// </summary>
public static class NvidiaBlackwellCommands
{
    public const uint CopyClassA = 0xC9B5;
    public const uint CopyClassB = 0xCAB5;
    public const int Fill32CommandBytes = 72;
    public const int GpfifoEntryBytes = 8;
    public const uint GpPutOffset = 0x8C;

    public static uint IncrementingMethodHeader(uint methodByteOffset, int subchannel, int valueCount)
    {
        if ((methodByteOffset & 3) != 0 || methodByteOffset > 0x3FFC || subchannel is < 0 or > 7 ||
            valueCount is < 1 or > 0x1FFF || (ulong)methodByteOffset + (ulong)(valueCount - 1) * 4 > 0x3FFC)
            throw new ArgumentOutOfRangeException(nameof(methodByteOffset), "Method fields or incremented method range do not fit.");
        return 1U << 29 | (uint)valueCount << 16 | (uint)subchannel << 13 | methodByteOffset >> 2;
    }

    /// <summary>
    /// Four-byte constant fill followed by a system-flushed release on the same
    /// copy engine. The fence is a 32-bit mapped word; it is not a frontend-only
    /// semaphore that could run ahead of asynchronous copy-engine work.
    /// </summary>
    public static int WriteFill32(Span<byte> destination, uint copyClass, NvidiaGpuRange target,
        uint pixelCount, uint color, NvidiaGpuRange fence, uint completionPayload)
    {
        if (copyClass != CopyClassA && copyClass != CopyClassB) throw new ArgumentOutOfRangeException(nameof(copyClass));
        target.RequireAlignment(4); fence.RequireAlignment(4);
        if (pixelCount == 0 || (ulong)pixelCount * 4 > target.Length || fence.Length < 4 || completionPayload == 0)
            throw new ArgumentOutOfRangeException(nameof(pixelCount));
        if (target.Overlaps(fence)) throw new ArgumentException("Fill target overlaps its completion fence.");
        if (destination.Length < Fill32CommandBytes) throw new ArgumentException("Pushbuffer is too short.", nameof(destination));

        var writer = new MethodWriter(destination);
        writer.Write(0x000, [copyClass]); // SET_OBJECT, fixed copy subchannel 4.
        writer.Write(0x704, [color, 5U | 3U << 16]); // CONST_B, one four-byte destination component.
        writer.Write(0x408, [(uint)(target.Address >> 32), (uint)target.Address]);
        writer.Write(0x418, [pixelCount]);
        // Non-pipelined; pitched source/destination; remap; system flush; PLC off;
        // virtual destination, default copy type, no scrub or physical aperture.
        writer.Write(0x300, [2U | 1U << 2 | 1U << 7 | 1U << 8 | 1U << 10 | 1U << 26]);
        writer.Write(0x240, [(uint)(fence.Address >> 32), (uint)fence.Address, completionPayload]);
        // No data transfer; system flush; one-word semaphore without timestamp.
        writer.Write(0x300, [1U << 2 | 1U << 3 | 1U << 26]);
        return writer.BytesWritten;
    }

    /// <summary>Set the 17 high pushbuffer-address bits before its ordinary entry.</summary>
    public static void WriteSegmentBase(Span<byte> destination, NvidiaGpuRange pushbuffer)
    {
        pushbuffer.RequireAlignment(4);
        if (destination.Length < GpfifoEntryBytes) throw new ArgumentException("Entry is too short.", nameof(destination));
        BinaryPrimitives.WriteUInt32LittleEndian(destination, (uint)(pushbuffer.Address >> 40) << 8);
        BinaryPrimitives.WriteUInt32LittleEndian(destination[4..], 4);
    }

    public static void WriteGpfifoEntry(Span<byte> destination, NvidiaGpuRange pushbuffer, int commandBytes)
    {
        pushbuffer.RequireAlignment(4);
        if (commandBytes <= 0 || (commandBytes & 3) != 0 || commandBytes / 4 > 0x1FFFFF ||
            (ulong)commandBytes > pushbuffer.Length ||
            (pushbuffer.Address >> 40) != ((pushbuffer.Address + (ulong)commandBytes - 1) >> 40))
            throw new ArgumentOutOfRangeException(nameof(commandBytes), "Pushbuffer length, alignment or 40-bit segment boundary is invalid.");
        if (destination.Length < GpfifoEntryBytes) throw new ArgumentException("Entry is too short.", nameof(destination));
        BinaryPrimitives.WriteUInt32LittleEndian(destination, (uint)pushbuffer.Address);
        // Main level, unconditional fetch, SYNC_WAIT. Higher address bits were
        // supplied by the preceding SET_PB_SEGMENT_EXTENDED_BASE control entry.
        BinaryPrimitives.WriteUInt32LittleEndian(destination[4..],
            (uint)(pushbuffer.Address >> 32) & 0xFF | (uint)(commandBytes / 4) << 10 | 1U << 31);
    }

    private ref struct MethodWriter
    {
        private readonly Span<byte> buffer;
        public MethodWriter(Span<byte> buffer) => this.buffer = buffer;
        public int BytesWritten { get; private set; }
        public void Write(uint method, scoped ReadOnlySpan<uint> values)
        {
            BinaryPrimitives.WriteUInt32LittleEndian(buffer[BytesWritten..], IncrementingMethodHeader(method, 4, values.Length));
            BytesWritten += 4;
            foreach (uint value in values)
            {
                BinaryPrimitives.WriteUInt32LittleEndian(buffer[BytesWritten..], value);
                BytesWritten += 4;
            }
        }
    }
}