XFESystem

using System.Diagnostics;
using System.Runtime.InteropServices;
using System.Text;
using System.Text.Json;

namespace XFESystem.DriverOracle;

internal sealed record CudaDevice(int Ordinal, int Handle, string Name, string PciBusId, string Uuid, ulong TotalMemoryBytes, int Major, int Minor, PciDevice? WindowsPciIdentity);
internal sealed record CudaCall(string Api, int ReturnCode, double Milliseconds, bool Cleanup);

internal sealed unsafe class CudaDriver : IDisposable
{
    private readonly nint library;
    private readonly StreamWriter trace;
    public string LibraryPath { get; }
    public List<CudaCall> Calls { get; } = [];
    public bool LibraryReleased { get; private set; }
    public int Version { get; private set; }

    [UnmanagedFunctionPointer(CallingConvention.Winapi)] internal delegate int InitD(uint flags);
    [UnmanagedFunctionPointer(CallingConvention.Winapi)] internal delegate int IntOutD(out int value);
    [UnmanagedFunctionPointer(CallingConvention.Winapi)] internal delegate int DeviceD(out int device, int ordinal);
    [UnmanagedFunctionPointer(CallingConvention.Winapi)] internal delegate int DeviceTextD(byte* text, int length, int device);
    [UnmanagedFunctionPointer(CallingConvention.Winapi)] internal delegate int UuidD(byte* uuid, int device);
    [UnmanagedFunctionPointer(CallingConvention.Winapi)] internal delegate int TotalMemD(out nuint bytes, int device);
    [UnmanagedFunctionPointer(CallingConvention.Winapi)] internal delegate int CapabilityD(out int major, out int minor, int device);
    [UnmanagedFunctionPointer(CallingConvention.Winapi)] internal delegate int ContextCreateD(out nint context, uint flags, int device);
    [UnmanagedFunctionPointer(CallingConvention.Winapi)] internal delegate int PointerD(nint value);
    [UnmanagedFunctionPointer(CallingConvention.Winapi)] internal delegate int PointerOutD(out nint value);
    [UnmanagedFunctionPointer(CallingConvention.Winapi)] internal delegate int NoArgsD();
    [UnmanagedFunctionPointer(CallingConvention.Winapi)] internal delegate int LinkCreateD(uint count, int* options, nint* values, out nint state);
    [UnmanagedFunctionPointer(CallingConvention.Winapi)] internal delegate int LinkAddD(nint state, int inputType, byte* data, nuint size, [MarshalAs(UnmanagedType.LPStr)] string name, uint count, int* options, nint* values);
    [UnmanagedFunctionPointer(CallingConvention.Winapi)] internal delegate int LinkCompleteD(nint state, out nint cubin, out nuint size);
    [UnmanagedFunctionPointer(CallingConvention.Winapi)] internal delegate int ModuleLoadD(out nint module, byte* image);
    [UnmanagedFunctionPointer(CallingConvention.Winapi)] internal delegate int FunctionD(out nint function, nint module, [MarshalAs(UnmanagedType.LPStr)] string name);
    [UnmanagedFunctionPointer(CallingConvention.Winapi)] internal delegate int FunctionAttributeD(out int value, int attribute, nint function);
    [UnmanagedFunctionPointer(CallingConvention.Winapi)] internal delegate int AllocateD(out ulong address, nuint bytes);
    [UnmanagedFunctionPointer(CallingConvention.Winapi)] internal delegate int FreeD(ulong address);
    [UnmanagedFunctionPointer(CallingConvention.Winapi)] internal delegate int UploadD(ulong target, void* source, nuint bytes);
    [UnmanagedFunctionPointer(CallingConvention.Winapi)] internal delegate int DownloadD(void* target, ulong source, nuint bytes);
    [UnmanagedFunctionPointer(CallingConvention.Winapi)] internal delegate int LaunchD(nint function, uint gx, uint gy, uint gz, uint bx, uint by, uint bz, uint sharedBytes, nint stream, void** arguments, void** extra);
    [DllImport("kernel32.dll", CharSet = CharSet.Unicode, SetLastError = true)]
    [DefaultDllImportSearchPaths(DllImportSearchPath.System32)]
    private static extern uint GetModuleFileNameW(nint module, StringBuilder path, int size);

    public CudaDriver(StreamWriter trace)
    {
        this.trace = trace;
        LibraryPath = Path.Combine(Environment.SystemDirectory, "nvcuda.dll");
        library = NativeLibrary.Load(LibraryPath);
        var actual = new StringBuilder(32768);
        uint count = GetModuleFileNameW(library, actual, actual.Capacity);
        if (count == 0 || count >= actual.Capacity || !Path.GetFullPath(actual.ToString()).Equals(Path.GetFullPath(LibraryPath), StringComparison.OrdinalIgnoreCase))
        {
            NativeLibrary.Free(library);
            throw new InvalidOperationException("The loaded CUDA driver did not resolve to the pinned System32 path.");
        }
    }

    internal T Api<T>(string name) where T : Delegate => Marshal.GetDelegateForFunctionPointer<T>(NativeLibrary.GetExport(library, name));
    internal void Check(string name, Func<int> call)
    {
        var timer = Stopwatch.StartNew(); int code = call();
        Record(name, code, timer.Elapsed.TotalMilliseconds, false);
        if (code != 0) throw new InvalidOperationException(name + " returned CUDA error " + code + ".");
    }
    // Pointer-bearing calls cannot be captured by a lambda; record immediately
    // after their return and before interpreting any output.
    internal void Check(string name, int code)
    {
        Record(name, code, 0, false);
        if (code != 0) throw new InvalidOperationException(name + " returned CUDA error " + code + ".");
    }
    internal bool Cleanup(string name, Func<int> call)
    {
        try { var timer = Stopwatch.StartNew(); int code = call(); Record(name, code, timer.Elapsed.TotalMilliseconds, true); return code == 0; }
        catch (Exception error) { trace.WriteLine(JsonSerializer.Serialize(new { Api = name, Cleanup = true, ManagedException = error.ToString() })); return false; }
    }
    private void Record(string name, int code, double elapsed, bool cleanup)
    {
        var entry = new CudaCall(name, code, elapsed, cleanup); Calls.Add(entry); trace.WriteLine(JsonSerializer.Serialize(entry));
    }

    public CudaDevice[] Query()
    {
        int version = 0; Check("cuDriverGetVersion", () => Api<IntOutD>("cuDriverGetVersion")(out version)); Version = version;
        Check("cuInit", () => Api<InitD>("cuInit")(0));
        int count = 0; Check("cuDeviceGetCount", () => Api<IntOutD>("cuDeviceGetCount")(out count));
        if (count is < 0 or > 32) throw new InvalidOperationException("CUDA device enumeration exceeds its bound.");
        var pciDevices = WindowsPci.Exact5090(); var devices = new List<CudaDevice>();
        byte* name = stackalloc byte[256]; byte* bus = stackalloc byte[32]; byte* uuid = stackalloc byte[16];
        for (int ordinal = 0; ordinal < count; ordinal++)
        {
            int device = 0; Check("cuDeviceGet", Api<DeviceD>("cuDeviceGet")(out device, ordinal));
            new Span<byte>(name, 256).Clear(); new Span<byte>(bus, 32).Clear();
            Check("cuDeviceGetName", Api<DeviceTextD>("cuDeviceGetName")(name, 256, device));
            Check("cuDeviceGetPCIBusId", Api<DeviceTextD>("cuDeviceGetPCIBusId")(bus, 32, device));
            Check("cuDeviceGetUuid", Api<UuidD>("cuDeviceGetUuid")(uuid, device));
            Check("cuDeviceTotalMem_v2", Api<TotalMemD>("cuDeviceTotalMem_v2")(out nuint memory, device));
            Check("cuDeviceComputeCapability", Api<CapabilityD>("cuDeviceComputeCapability")(out int major, out int minor, device));
            string deviceName = Marshal.PtrToStringUTF8((nint)name) ?? "", busId = Marshal.PtrToStringUTF8((nint)bus) ?? "";
            var matched = pciDevices.Where(pci => busId.Equals($"0000:{pci.Bus:x2}:{pci.Device:x2}.{pci.Function:x}", StringComparison.OrdinalIgnoreCase)).ToArray();
            if (matched.Length > 1) throw new InvalidOperationException("The CUDA bus has ambiguous Windows PCI identities.");
            devices.Add(new(ordinal, device, deviceName, busId, Convert.ToHexStringLower(new ReadOnlySpan<byte>(uuid, 16)), memory, major, minor, matched.SingleOrDefault()));
        }
        return devices.ToArray();
    }

    public void Dispose()
    {
        if (!LibraryReleased) { NativeLibrary.Free(library); LibraryReleased = true; }
    }
}