XFESystem

namespace XFESystem.Gpu.Nvidia;

/// <summary>
/// Publishes real Blackwell copy commands into caller-owned native mappings.
/// Prerequisites: GSP/RM initialized the device/channel, GPU MMU mapped these
/// ranges, and RM supplied GPPut and the work-submission register/token. This
/// class does not perform PCI discovery, map BARs, or invent a doorbell offset.
/// The current XFESystem runtime has no such resource provider and does not
/// construct this class. Host tests use private unmanaged memory, not a GPU.
/// </summary>
public sealed unsafe class NvidiaMappedCopyChannel : IDisposable
{
    private readonly object gate = new();
    private readonly NvidiaMappedRegion commands, fifo, fence;
    private readonly nuint gpPut, workSubmission, drainBar1;
    private readonly uint workToken, copyClass;
    private readonly int entryCount;
    private int producer;
    private uint payload;
    private bool inFlight, disposed;

    public bool InFlight => Volatile.Read(ref inFlight);

    public NvidiaMappedCopyChannel(uint copyClass, NvidiaMappedRegion commands, NvidiaMappedRegion fifo,
        int entryCount, NvidiaMappedRegion fence, nuint gpPut, nuint workSubmission,
        uint workToken, int initialProducer = 0, nuint drainBar1 = 0)
    {
        if (copyClass != NvidiaBlackwellCommands.CopyClassA && copyClass != NvidiaBlackwellCommands.CopyClassB)
            throw new ArgumentOutOfRangeException(nameof(copyClass));
        if (commands.Length < NvidiaBlackwellCommands.Fill32CommandBytes || entryCount < 4 ||
            entryCount > 65536 || (entryCount & (entryCount - 1)) != 0 ||
            fifo.Length < (long)entryCount * NvidiaBlackwellCommands.GpfifoEntryBytes || fence.Length != 4 ||
            initialProducer < 0 || initialProducer >= entryCount)
            throw new ArgumentException("Invalid command, FIFO, fence or producer geometry.");
        fifo.Gpu.RequireAlignment(8);
        if ((fifo.CpuAddress & 7) != 0) throw new ArgumentException("GPFIFO CPU mapping must be eight-byte aligned.");
        if (commands.Gpu.Overlaps(fifo.Gpu) || commands.Gpu.Overlaps(fence.Gpu) || fifo.Gpu.Overlaps(fence.Gpu) ||
            commands.CpuOverlaps(fifo) || commands.CpuOverlaps(fence) || fifo.CpuOverlaps(fence))
            throw new ArgumentException("Channel memory ranges overlap.");
        RequireRegister(gpPut); RequireRegister(workSubmission);
        if (drainBar1 != 0) RequireRegister(drainBar1);
        if (gpPut == workSubmission) throw new ArgumentException("GPPut and the work-submission register alias.");
        foreach (var memory in new[] { commands, fifo, fence })
            if (Contains(memory, gpPut) || Contains(memory, workSubmission) || (drainBar1 != 0 && Contains(memory, drainBar1)))
                throw new ArgumentException("A control register aliases command, FIFO or fence memory.");
        // Validate the complete pushbuffer range before exposing an instance.
        Span<byte> scratch = stackalloc byte[NvidiaBlackwellCommands.GpfifoEntryBytes];
        NvidiaBlackwellCommands.WriteGpfifoEntry(scratch, commands.Gpu, NvidiaBlackwellCommands.Fill32CommandBytes);
        this.copyClass = copyClass; this.commands = commands; this.fifo = fifo; this.fence = fence;
        this.entryCount = entryCount; this.gpPut = gpPut; this.workSubmission = workSubmission;
        this.workToken = workToken; this.drainBar1 = drainBar1; producer = initialProducer;
    }

    public uint SubmitFill32(NvidiaGpuRange target, uint pixels, uint color)
    {
        Monitor.Enter(gate);
        uint result;
        // The native runtime does not unwind arbitrary finally handlers. Release
        // explicitly on both paths; compiler-generated lock/finally is unsafe here.
        try { result = SubmitFill32Core(target, pixels, color); }
        catch { Monitor.Exit(gate); throw; }
        Monitor.Exit(gate);
        return result;
    }

    private uint SubmitFill32Core(NvidiaGpuRange target, uint pixels, uint color)
    {
            ObjectDisposedException.ThrowIf(disposed, this);
            if (inFlight) throw new InvalidOperationException("The prior GPU fence has not completed; its mappings cannot be reused.");
            if (payload == uint.MaxValue) throw new InvalidOperationException("Fence payload space is exhausted; recreate the channel after draining it.");
            if (target.Overlaps(commands.Gpu) || target.Overlaps(fifo.Gpu) || target.Overlaps(fence.Gpu))
                throw new ArgumentException("Fill target aliases live channel memory.");
            uint next = payload + 1;
            // Validate/encode privately first. Failed validation leaves all DMA
            // mappings and registers untouched, including the previous fence.
            Span<byte> packet = stackalloc byte[NvidiaBlackwellCommands.Fill32CommandBytes];
            NvidiaBlackwellCommands.WriteFill32(packet, copyClass, target, pixels, color, fence.Gpu, next);
            Span<byte> entries = stackalloc byte[NvidiaBlackwellCommands.GpfifoEntryBytes * 2];
            NvidiaBlackwellCommands.WriteSegmentBase(entries, commands.Gpu);
            NvidiaBlackwellCommands.WriteGpfifoEntry(entries[8..], commands.Gpu, packet.Length);
            // From here onwards, conservatively retain ownership even if a
            // mapping faults while publishing. Never recycle a partial submit.
            payload = next; inFlight = true;
            Volatile.Write(ref *(uint*)fence.CpuAddress, 0);
            packet.CopyTo(new Span<byte>((void*)commands.CpuAddress, packet.Length));
            int second = (producer + 1) & (entryCount - 1);
            entries[..8].CopyTo(new Span<byte>((void*)(fifo.CpuAddress + (nuint)(producer * 8)), 8));
            entries[8..].CopyTo(new Span<byte>((void*)(fifo.CpuAddress + (nuint)(second * 8)), 8));
            Thread.MemoryBarrier();
            if (drainBar1 != 0) _ = Volatile.Read(ref *(uint*)drainBar1);
            producer = (producer + 2) & (entryCount - 1);
            Volatile.Write(ref *(uint*)gpPut, (uint)producer);
            Thread.MemoryBarrier();
            Volatile.Write(ref *(uint*)workSubmission, workToken);
            return payload;
    }

    /// <summary>
    /// Reads the mapped GPU-written fence. A timeout never implies completion and no
    /// CPU-side simulator writes this word. Callers must retain mappings on timeout.
    /// </summary>
    public bool TryRetire()
    {
        Monitor.Enter(gate);
        bool result;
        try { result = TryRetireCore(); }
        catch { Monitor.Exit(gate); throw; }
        Monitor.Exit(gate);
        return result;
    }

    private bool TryRetireCore()
    {
            ObjectDisposedException.ThrowIf(disposed, this);
            if (!inFlight) return true;
            if (Volatile.Read(ref *(uint*)fence.CpuAddress) != payload) return false;
            Thread.MemoryBarrier(); inFlight = false;
            return true;
    }

    public void Dispose()
    {
        Monitor.Enter(gate);
        try { DisposeCore(); }
        catch { Monitor.Exit(gate); throw; }
        Monitor.Exit(gate);
    }

    private void DisposeCore()
    {
            if (disposed) return;
            if (inFlight) throw new InvalidOperationException("Wait for the GPU fence or reset the owning channel before unmapping memory.");
            disposed = true;
    }

    private static bool Contains(NvidiaMappedRegion memory, nuint word) =>
        word < memory.CpuAddress + (nuint)memory.Length && memory.CpuAddress < word + 4;
    private static void RequireRegister(nuint word)
    {
        if (word == 0 || (word & 3) != 0 || word > nuint.MaxValue - 4)
            throw new ArgumentException("A mapped control register is null, unaligned or overflows.");
    }
}