XFESystem

using System.Buffers.Binary;
using System.Runtime.InteropServices;
using System.Reflection;
using XFESystem.Gpu.Nvidia;

namespace XFESystem.Gpu.Nvidia.Tests;

public sealed class NvidiaProtocolTests
{
    [Theory]
    [InlineData(0x000U, 4, 1, 0x20018000U)]
    [InlineData(0x704U, 4, 2, 0x200281C1U)]
    [InlineData(0x3FFCU, 7, 1, 0x2001EFFFU)]
    public void MethodHeadersUsePublishedBitFields(uint method, int subchannel, int count, uint expected) =>
        Assert.Equal(expected, NvidiaBlackwellCommands.IncrementingMethodHeader(method, subchannel, count));

    [Theory]
    [InlineData(1U, 4, 1)]
    [InlineData(0x4000U, 4, 1)]
    [InlineData(0U, 8, 1)]
    [InlineData(0U, -1, 1)]
    [InlineData(0U, 4, 0)]
    [InlineData(0U, 4, 8192)]
    [InlineData(0x3FFCU, 4, 2)]
    public void InvalidMethodFieldsAreRejected(uint method, int subchannel, int count) =>
        Assert.Throws<ArgumentOutOfRangeException>(() => NvidiaBlackwellCommands.IncrementingMethodHeader(method, subchannel, count));

    [Theory]
    [InlineData(NvidiaBlackwellCommands.CopyClassA)]
    [InlineData(NvidiaBlackwellCommands.CopyClassB)]
    public void FillAndSystemReleaseAreExactLittleEndianCommands(uint copyClass)
    {
        var bytes = new byte[72];
        NvidiaBlackwellCommands.WriteFill32(bytes, copyClass, new(0x123456789ABC, 8), 2, 0xFFAC99EA,
            new(0x0111222233334400, 4), 0x98765432);
        uint[] words = [
            0x20018000, copyClass,
            0x200281C1, 0xFFAC99EA, 0x00030005,
            0x20028102, 0x00001234, 0x56789ABC,
            0x20018106, 2,
            0x200180C0, 0x04000586,
            0x20038090, 0x01112222, 0x33334400, 0x98765432,
            0x200180C0, 0x0400000C
        ];
        Assert.Equal(words, Enumerable.Range(0, 18).Select(index => BinaryPrimitives.ReadUInt32LittleEndian(bytes.AsSpan(index * 4))));
        Assert.Equal(new byte[] { 0, 0x80, 1, 0x20 }, bytes[..4]);
    }

    [Fact]
    public void SegmentBaseAndGpfifoEntryRetainTheWhole57BitAddress()
    {
        var push = new NvidiaGpuRange(0x0123456789ABC000, 72);
        var bytes = new byte[16];
        NvidiaBlackwellCommands.WriteSegmentBase(bytes, push);
        NvidiaBlackwellCommands.WriteGpfifoEntry(bytes.AsSpan(8), push, 72);
        Assert.Equal("004523010400000000C0AB8967480080", Convert.ToHexString(bytes));
    }

    [Theory]
    [InlineData(0UL, 4UL)]
    [InlineData(1UL, 0UL)]
    [InlineData(ulong.MaxValue, 4UL)]
    [InlineData(NvidiaGpuRange.AddressLimit - 4, 8UL)]
    [InlineData(NvidiaGpuRange.AddressLimit, 4UL)]
    public void EmptyOverflowingOrTooWideGpuRangesAreRejected(ulong address, ulong length) =>
        Assert.Throws<ArgumentOutOfRangeException>(() => new NvidiaGpuRange(address, length));

    [Fact]
    public void InvalidPacketsDoNotPartiallyModifyCallerStorage()
    {
        byte[] bytes = Enumerable.Repeat((byte)0xCD, 80).ToArray();
        var target = new NvidiaGpuRange(0x10000, 16);
        var fence = new NvidiaGpuRange(0x20000, 4);
        Assert.Throws<ArgumentException>(() => NvidiaBlackwellCommands.WriteFill32(bytes.AsSpan(0, 71), 0xCAB5, target, 4, 0, fence, 1));
        Assert.Throws<ArgumentOutOfRangeException>(() => NvidiaBlackwellCommands.WriteFill32(bytes, 0xCEC0, target, 4, 0, fence, 1));
        Assert.Throws<ArgumentOutOfRangeException>(() => NvidiaBlackwellCommands.WriteFill32(bytes, 0xCAB5, target, 5, 0, fence, 1));
        Assert.Throws<ArgumentException>(() => NvidiaBlackwellCommands.WriteFill32(bytes, 0xCAB5, target, 4, 0, target, 1));
        Assert.All(bytes, value => Assert.Equal((byte)0xCD, value));
    }

    [Fact]
    public void GpfifoRejectsSegmentCrossingTruncationAndUnalignedLengths()
    {
        var bytes = new byte[8];
        Assert.Throws<ArgumentOutOfRangeException>(() => NvidiaBlackwellCommands.WriteGpfifoEntry(bytes, new((1UL << 40) - 4, 72), 72));
        Assert.Throws<ArgumentOutOfRangeException>(() => NvidiaBlackwellCommands.WriteGpfifoEntry(bytes, new(0x1000, 72), 73));
        Assert.Throws<ArgumentOutOfRangeException>(() => NvidiaBlackwellCommands.WriteGpfifoEntry(bytes, new(0x1000, 72), 76));
        Assert.Throws<ArgumentOutOfRangeException>(() => NvidiaBlackwellCommands.WriteGpfifoEntry(bytes, new(0x1000, 1UL << 24), 1 << 23));
        Assert.Equal(new byte[8], bytes);
    }

    [Fact]
    public unsafe void PrivateUnmanagedMemoryPublicationProtectsInFlightLifetimeAndWrapsFifo()
    {
        using var arena = new PrivateArena();
        using var channel = arena.Create(initialProducer: 3);
        uint first = channel.SubmitFill32(new(0x40000, 16), 4, 0xFFAC99EA);
        Assert.Equal(1U, first);
        Assert.True(channel.InFlight);
        Assert.Equal(1U, arena.Read(320));
        Assert.Equal(0xBAADF00DU, arena.Read(324));
        Assert.Equal(0U, arena.Read(256));
        Assert.Equal(0x00452300U, arena.Read(128 + 3 * 8));
        Assert.Equal(4U, arena.Read(128 + 3 * 8 + 4));
        Assert.Equal(0x80004800U, arena.Read(128 + 4));
        Assert.False(channel.TryRetire());
        Assert.Throws<InvalidOperationException>(() => channel.SubmitFill32(new(0x40000, 16), 4, 0));
        Assert.Throws<InvalidOperationException>(() => channel.Dispose());
        // This explicit CPU write is solely an offline fence-observation test.
        // It is not a GPU result and does not validate hardware command execution.
        arena.Write(256, first + 1);
        Assert.False(channel.TryRetire());
        arena.Write(256, first);
        Assert.True(channel.TryRetire());
        Assert.False(channel.InFlight);
        uint second = channel.SubmitFill32(new(0x40000, 16), 4, 0x12345678);
        Assert.Equal(2U, second);
        Assert.Equal(3U, arena.Read(320));
        Assert.Equal(0U, arena.Read(256));
        arena.Write(256, second);
        Assert.True(channel.TryRetire());
    }

    [Fact]
    public void InvalidSubmissionLeavesEveryMappingAndRegisterUnmodified()
    {
        using var arena = new PrivateArena();
        using var channel = arena.Create();
        byte[] before = arena.Snapshot();
        Assert.Throws<ArgumentOutOfRangeException>(() => channel.SubmitFill32(new(0x40000, 16), 5, 0));
        Assert.Throws<ArgumentException>(() => channel.SubmitFill32(new(0x30000, 4), 1, 0));
        Assert.Equal(before, arena.Snapshot());
        Assert.False(channel.InFlight);
    }

    [Fact]
    public void CpuGpuAndControlAliasesAreRejectedBeforeAnyWrite()
    {
        using var arena = new PrivateArena();
        var command = arena.Region(0, 0x0045230000000000, 72);
        var fifo = arena.Region(128, 0x20000, 32);
        var fence = arena.Region(256, 0x30000, 4);
        byte[] before = arena.Snapshot();
        Assert.Throws<ArgumentException>(() => new NvidiaMappedCopyChannel(0xCAB5, command, fifo, 4,
            arena.Region(0, 0x30000, 4), arena.Pointer(320), arena.Pointer(324), 1));
        Assert.Throws<ArgumentException>(() => new NvidiaMappedCopyChannel(0xCAB5, command, fifo, 4,
            arena.Region(256, 0x20000, 4), arena.Pointer(320), arena.Pointer(324), 1));
        Assert.Throws<ArgumentException>(() => new NvidiaMappedCopyChannel(0xCAB5, command, fifo, 4,
            fence, arena.Pointer(0), arena.Pointer(324), 1));
        Assert.Equal(before, arena.Snapshot());
    }

    [Fact]
    public void DisposedChannelCannotPublishOrReadFreedMappings()
    {
        using var arena = new PrivateArena();
        var channel = arena.Create();
        channel.Dispose(); channel.Dispose();
        Assert.Throws<ObjectDisposedException>(() => channel.SubmitFill32(new(0x40000, 16), 4, 0));
        Assert.Throws<ObjectDisposedException>(() => channel.TryRetire());
    }

    [Fact]
    public async Task ValidationAndInFlightExceptionsReleaseTheGateForAnotherThread()
    {
        using var arena = new PrivateArena();
        using var channel = arena.Create();
        Assert.Throws<ArgumentOutOfRangeException>(() => channel.SubmitFill32(new(0x40000, 16), 5, 0));
        Assert.True(await Task.Run(channel.TryRetire).WaitAsync(TimeSpan.FromSeconds(2)));
        uint payload = channel.SubmitFill32(new(0x40000, 16), 4, 0);
        Assert.Throws<InvalidOperationException>(() => channel.Dispose());
        Assert.False(await Task.Run(channel.TryRetire).WaitAsync(TimeSpan.FromSeconds(2)));
        arena.Write(256, payload);
        Assert.True(await Task.Run(channel.TryRetire).WaitAsync(TimeSpan.FromSeconds(2)));
    }

    [Fact]
    public void NativeChannelGateDoesNotDependOnUnsupportedFinallyUnwind()
    {
        foreach (var method in typeof(NvidiaMappedCopyChannel).GetMethods(BindingFlags.Public | BindingFlags.NonPublic | BindingFlags.Instance | BindingFlags.DeclaredOnly))
            if (method.GetMethodBody() is { } body)
                Assert.DoesNotContain(body.ExceptionHandlingClauses, clause => clause.Flags == ExceptionHandlingClauseOptions.Finally);
    }

    private sealed unsafe class PrivateArena : IDisposable
    {
        private readonly nint allocation = Marshal.AllocHGlobal(512);
        public PrivateArena() => new Span<byte>((void*)allocation, 512).Fill(0xCD);
        public nuint Pointer(int offset) => (nuint)(allocation + offset);
        public NvidiaMappedRegion Region(int offset, ulong gpu, int length) => new(Pointer(offset), new(gpu, (ulong)length), length);
        public uint Read(int offset) => *(uint*)(allocation + offset);
        public void Write(int offset, uint value) => *(uint*)(allocation + offset) = value;
        public byte[] Snapshot() => new Span<byte>((void*)allocation, 512).ToArray();
        public NvidiaMappedCopyChannel Create(int initialProducer = 0) => new(0xCAB5,
            Region(0, 0x0045230000000000, 72), Region(128, 0x20000, 32), 4,
            Region(256, 0x30000, 4), Pointer(320), Pointer(324), 0xBAADF00D, initialProducer, Pointer(328));
        public void Dispose() => Marshal.FreeHGlobal(allocation);
    }
}