XFESystem

using XFESystem.Boot;

namespace XFESystem.NativeKernel;

// Read the immutable, pointer-free PCI resource snapshot published by our UEFI
// shim. UEFI Boot/Runtime Services are never called after ExitBootServices.
public static unsafe class NativeGpuFirmware
{
    static FirmwareGpuHandoff* _report;
    static bool _initialized;

    static byte* PhysicalPointer(NativeBootInfo* boot, ulong physical, ulong bytes)
    {
        if (boot == null || boot->MemoryEntries == null || boot->MemoryEntryCount == 0 || boot->MemoryEntryCount > 4096 ||
            physical == 0 || bytes == 0 || physical >= (1UL << 52) || bytes > (1UL << 52) - physical ||
            boot->Hhdm == 0 || physical > ~0UL - boot->Hhdm || bytes - 1 > ~0UL - (boot->Hhdm + physical)) return null;
        // EfiRuntimeServicesData and the EFI table are Reserved(Mapped) under
        // Limine base6. Never dereference an address merely because it is HHDM.
        bool covered = false;
        for (ulong i = 0; i < boot->MemoryEntryCount; i++)
        {
            MemoryMapEntry* entry = boot->MemoryEntries[i];
            if (entry == null || entry->Length > ~0UL - entry->Base) return null;
            if (entry->Length != 0 && physical < entry->Base + entry->Length && entry->Base < physical + bytes && entry->Kind != 8)
                return null;
            if (entry != null && entry->Kind == 8 && physical >= entry->Base &&
                physical - entry->Base < entry->Length && bytes <= entry->Length - (physical - entry->Base))
                covered = true;
        }
        return covered ? (byte*)(boot->Hhdm + physical) : null;
    }

    static uint HeaderCrc(byte* header, uint bytes)
    {
        uint crc = 0xffffffff;
        for (uint i = 0; i < bytes; i++)
        {
            crc ^= i >= 16 && i < 20 ? (byte)0 : header[i];
            for (uint bit = 0; bit < 8; bit++) crc = (crc >> 1) ^ ((crc & 1) != 0 ? 0xedb88320u : 0);
        }
        return ~crc;
    }

    public static void Initialize(NativeBootInfo* boot)
    {
        if (_initialized) return;
        _initialized = true;
        _report = Locate(boot);
        Serial.Write("XFE_GPU_FIRMWARE_HANDOFF|Valid|"); Serial.Number(_report == null ? 0UL : 1UL);
        Serial.Write("|Count|"); Serial.Number(_report == null ? 0UL : _report->Count);
        Serial.Line("|ResourceOnly|1|Cuda|0|GpuRasterization|0");
    }

    static FirmwareGpuHandoff* Locate(NativeBootInfo* boot)
    {
        if (boot == null || boot->MemoryEntries == null || boot->MemoryEntryCount == 0 ||
            boot->MemoryEntryCount > 4096 || boot->EfiSystemTable < boot->Hhdm) return null;
        ulong physical = boot->EfiSystemTable - boot->Hhdm;
        byte* table = PhysicalPointer(boot, physical, 120);
        if (table == null || *(ulong*)table != 0x5453595320494249UL || *(uint*)(table + 20) != 0) return null;
        uint headerBytes = *(uint*)(table + 12);
        if (headerBytes < 120 || headerBytes > 4096 || PhysicalPointer(boot, physical, headerBytes) == null ||
            *(uint*)(table + 16) != HeaderCrc(table, headerBytes)) return null;
        ulong count = *(ulong*)(table + 104), arrayPhysical = *(ulong*)(table + 112);
        if (count == 0 || count > 4096) return null;
        byte* entries = PhysicalPointer(boot, arrayPhysical, count * 24);
        if (entries == null) return null;
        FirmwareGpuHandoff* selected = null;
        for (ulong i = 0; i < count; i++)
        {
            byte* entry = entries + i * 24;
            if (*(ulong*)entry != FirmwareGpuHandoff.GuidLow || *(ulong*)(entry + 8) != FirmwareGpuHandoff.GuidHigh) continue;
            FirmwareGpuHandoff* report = (FirmwareGpuHandoff*)PhysicalPointer(boot, *(ulong*)(entry + 16), FirmwareGpuHandoff.TotalBytes);
            if (selected != null || !FirmwareGpuResources.Validate(report, FirmwareGpuHandoff.TotalBytes) || report->Flags != 0) return null;
            selected = report;
        }
        return selected;
    }

    public static void Apply(NativeBootInfo* boot, NativeGpuDeviceDiscovery* record)
    {
        Initialize(boot);
        if (_report == null || record == null) return;
        FirmwareGpuDevice* selected = null;
        for (uint i = 0; i < _report->Count; i++)
        {
            FirmwareGpuDevice* item = (FirmwareGpuDevice*)(_report->Devices + i * FirmwareGpuHandoff.DeviceBytes);
            if (item->Segment != record->Segment || item->Bus != record->Bus || item->Device != record->Device ||
                item->Function != record->Function) continue;
            if (selected != null || item->VendorId != record->VendorId || item->DeviceId != record->DeviceId ||
                item->ClassCode != record->ClassCode || (item->Command & 3) != (record->Command & 3)) return;
            selected = item;
        }
        if (selected == null) return;
        for (uint i = 0; i < 6; i++)
        {
            FirmwareGpuBar* source = (FirmwareGpuBar*)(selected->Bars + i * 24);
            if (source->Index != i || (source->Flags & 7) != 3 || (source->Flags & ~31u) != 0) continue;
            NativeGpuBarDiscovery* target = &record->Bar0 + i;
            // Firmware descriptors cannot reinterpret an I/O BAR or the high
            // DWORD of a 64-bit BAR. ApplyResourceSize checks the live address.
            if ((target->Flags & 1) == 0 || ((target->Flags & 4) != 0) != ((source->Flags & 8) != 0) ||
                ((target->Flags & 8) != 0) != ((source->Flags & 16) != 0)) continue;
            GpuHardwareBridge.ApplyResourceSize(target, source->Address, source->Bytes, 2);
        }
    }
}