// PCI bus driver. This file owns PCI enumeration, PDO creation, BAR decoding, and explicit MMIO mappings for function drivers.

#include "pci.h"

#include <stdint.h>
#include <stdbool.h>
#include <stddef.h>
#include <string.h>
#include <common/printf.h>

#include <ntlite/ddi/hal_api.h>
#include <ntlite/ddi/io_driver.h>
#include <kernel/mm_virt.h>
#include <kernel/ke_debug.h>
#include <kernel/ke_irp.h>
#include <kernel/ke_irq.h>
#include <kernel/hal/hal_apic.h>

int _fltused; // MSVC floating-point helper anchor required by printf.obj

static const HAL_API* gHal;
static DRIVER_OBJECT  gPciDriver;
static uint32_t       gPciDeviceCount;

// PCI MMIO BARs use a private virtual aperture instead of the physical
// direct map.  This avoids cacheability aliasing with large direct-map pages
// and keeps BAR mappings page-granular.

#define PCI_MMIO_VA_BASE  0xFFFF900000000000ULL
#define PCI_MMIO_VA_LIMIT 0xFFFF980000000000ULL

static uintptr_t gPciMmioNext = PCI_MMIO_VA_BASE;

// Implements the align Up64 operation.
static uint64_t PciAlignUp64(uint64_t value, uint64_t alignment) {
    return (value + alignment - 1u) & ~(alignment - 1u);
}

// Implements the default Mmio Flags operation.
static uint64_t PciDefaultMmioFlags(void) {
    return PAGE_PRESENT | PAGE_WRITE | PAGE_WRITETHROUGH | PAGE_NOCACHE | PAGE_NX;
}

// Creates mmio Va.
static KSTATUS PciAllocateMmioVa(uint64_t length, uintptr_t* outVa) {
    if (!outVa || !length)
        return STATUS_INVALID_PARAMETER;

    uint64_t alignedLength = PciAlignUp64(length, PAGE_SIZE);
    uintptr_t va = (uintptr_t)PciAlignUp64((uint64_t)gPciMmioNext, PAGE_SIZE);
    if (va < PCI_MMIO_VA_BASE || va >= PCI_MMIO_VA_LIMIT)
        return STATUS_NO_MEMORY;
    if (alignedLength > (uint64_t)(PCI_MMIO_VA_LIMIT - va))
        return STATUS_NO_MEMORY;

    gPciMmioNext = va + (uintptr_t)alignedLength;
    *outVa = va;
    return STATUS_SUCCESS;
}

// Checks whether pdo.
static bool PciIsPdo(DEVICE_OBJECT* pdo, PCI_PDO_EXTENSION** outExt) {
    if (!pdo || pdo->Type != DeviceObjectPhysical || pdo->u.Pdo.BusType != IoBusPci)
        return false;
    PCI_PDO_EXTENSION* ext = (PCI_PDO_EXTENSION*)pdo->DeviceExtension;
    if (!ext) return false;
    if (outExt) *outExt = ext;
    return true;
}

// Implements the raw Read32 operation.
static uint32_t PciRawRead32(uint8_t bus, uint8_t dev, uint8_t fn, uint16_t offset) {
    if (!gHal || offset > 0xFCu) return 0xFFFFFFFFu;
    return gHal->Pci.Read32(0, bus, dev, fn, offset & 0xFCu);
}

// Implements the raw Write32 operation.
static void PciRawWrite32(uint8_t bus, uint8_t dev, uint8_t fn, uint16_t offset, uint32_t value) {
    if (!gHal || offset > 0xFCu) return;
    gHal->Pci.Write32(0, bus, dev, fn, offset & 0xFCu, value);
}

// Implements the raw Read8 operation.
static uint8_t PciRawRead8(uint8_t bus, uint8_t dev, uint8_t fn, uint16_t offset) {
    uint32_t v = PciRawRead32(bus, dev, fn, offset);
    return (uint8_t)(v >> ((offset & 3u) * 8u));
}

// Implements the raw Read16 operation.
static uint16_t PciRawRead16(uint8_t bus, uint8_t dev, uint8_t fn, uint16_t offset) {
    uint32_t v = PciRawRead32(bus, dev, fn, offset);
    return (uint16_t)(v >> ((offset & 2u) * 8u));
}

// Implements the raw Write8 operation.
static void PciRawWrite8(uint8_t bus, uint8_t dev, uint8_t fn, uint16_t offset, uint8_t value) {
    uint16_t aligned = offset & ~3u;
    uint32_t shift = (offset & 3u) * 8u;
    uint32_t old = PciRawRead32(bus, dev, fn, aligned);
    old &= ~(0xFFu << shift);
    old |= ((uint32_t)value << shift);
    PciRawWrite32(bus, dev, fn, aligned, old);
}

// Implements the raw Write16 operation.
static void PciRawWrite16(uint8_t bus, uint8_t dev, uint8_t fn, uint16_t offset, uint16_t value) {
    uint16_t aligned = offset & ~3u;
    uint32_t shift = (offset & 2u) * 8u;
    uint32_t old = PciRawRead32(bus, dev, fn, aligned);
    old &= ~(0xFFFFu << shift);
    old |= ((uint32_t)value << shift);
    PciRawWrite32(bus, dev, fn, aligned, old);
}

// Implements the bar Size From Mask operation.
static uint64_t PciBarSizeFromMask(uint64_t mask, uint64_t baseMask) {
    uint64_t sizeMask = mask & baseMask;
    if (!sizeMask || sizeMask == baseMask)
        return 0;
    return (~sizeMask + 1u) & baseMask;
}

// Implements the header Bar Count operation.
static uint8_t PciHeaderBarCount(uint8_t headerType) {
    switch (headerType & PCI_HEADER_TYPE_MASK) {
        case 0x00: return 6;
        case 0x01: return 2;
        case 0x02: return 1;
        default:   return 0;
    }
}

// Implements the decode Bars operation.
static void PciDecodeBars(PCI_PDO_EXTENSION* ext) {
    uint8_t barCount = PciHeaderBarCount(ext->Info.HeaderType);
    for (uint8_t i = 0; i < PCI_MAX_BARS; i++) {
        ext->Bars[i].Type = PciBarUnused;
        ext->Bars[i].Index = i;
    }

    // Size BARs with I/O and memory decoding disabled so all-ones probes do
    // not accidentally alias a live resource.  Restore policy immediately;
    // PciEnableDevice() is the only exported path that changes command bits.
    uint16_t savedCommand = PciRawRead16(ext->Info.Address.Bus,
                                         ext->Info.Address.Device,
                                         ext->Info.Address.Function,
                                         PCI_REG_COMMAND_STATUS);
    uint16_t disabledCommand = (uint16_t)(savedCommand &
        (uint16_t)~(PCI_COMMAND_IO_SPACE | PCI_COMMAND_MEMORY_SPACE));
    if (disabledCommand != savedCommand) {
        PciRawWrite16(ext->Info.Address.Bus,
                      ext->Info.Address.Device,
                      ext->Info.Address.Function,
                      PCI_REG_COMMAND_STATUS,
                      disabledCommand);
    }

    for (uint8_t i = 0; i < barCount && i < PCI_MAX_BARS; i++) {
        uint16_t reg = (uint16_t)(PCI_REG_BAR0 + i * 4u);
        uint32_t original = PciRawRead32(ext->Info.Address.Bus,
                                         ext->Info.Address.Device,
                                         ext->Info.Address.Function,
                                         reg);
        if (original == 0 || original == 0xFFFFFFFFu)
            continue;

        PCI_BAR* bar = &ext->Bars[i];
        bar->Index = i;
        bar->Length = 0;
        bar->Prefetchable = false;

        if (original & 0x1u) {
            PciRawWrite32(ext->Info.Address.Bus,
                          ext->Info.Address.Device,
                          ext->Info.Address.Function,
                          reg,
                          0xFFFFFFFFu);
            uint32_t maskLo = PciRawRead32(ext->Info.Address.Bus,
                                           ext->Info.Address.Device,
                                           ext->Info.Address.Function,
                                           reg);
            PciRawWrite32(ext->Info.Address.Bus,
                          ext->Info.Address.Device,
                          ext->Info.Address.Function,
                          reg,
                          original);

            bar->Type = PciBarIo;
            bar->Base = (uint64_t)(original & ~0x3u);
            bar->Length = PciBarSizeFromMask((uint64_t)maskLo, 0xFFFFFFFCull);
            continue;
        }

        uint32_t memType = (original >> 1) & 0x3u;
        bar->Prefetchable = (original & 0x8u) != 0;

        if (memType == 0x2u && i + 1u < barCount && i + 1u < PCI_MAX_BARS) {
            uint16_t regHi = (uint16_t)(reg + 4u);
            uint32_t originalHi = PciRawRead32(ext->Info.Address.Bus,
                                               ext->Info.Address.Device,
                                               ext->Info.Address.Function,
                                               regHi);

            PciRawWrite32(ext->Info.Address.Bus,
                          ext->Info.Address.Device,
                          ext->Info.Address.Function,
                          reg,
                          0xFFFFFFFFu);
            PciRawWrite32(ext->Info.Address.Bus,
                          ext->Info.Address.Device,
                          ext->Info.Address.Function,
                          regHi,
                          0xFFFFFFFFu);
            uint32_t maskLo = PciRawRead32(ext->Info.Address.Bus,
                                           ext->Info.Address.Device,
                                           ext->Info.Address.Function,
                                           reg);
            uint32_t maskHi = PciRawRead32(ext->Info.Address.Bus,
                                           ext->Info.Address.Device,
                                           ext->Info.Address.Function,
                                           regHi);
            PciRawWrite32(ext->Info.Address.Bus,
                          ext->Info.Address.Device,
                          ext->Info.Address.Function,
                          reg,
                          original);
            PciRawWrite32(ext->Info.Address.Bus,
                          ext->Info.Address.Device,
                          ext->Info.Address.Function,
                          regHi,
                          originalHi);

            bar->Type = PciBarMemory64;
            bar->Base = ((uint64_t)originalHi << 32) | (uint64_t)(original & ~0xFu);
            uint64_t mask = ((uint64_t)maskHi << 32) | (uint64_t)(maskLo & ~0xFu);
            bar->Length = PciBarSizeFromMask(mask, 0xFFFFFFFFFFFFFFF0ull);

            i++;
            ext->Bars[i].Type = PciBarUnused;
            ext->Bars[i].Index = i;
        } else {
            PciRawWrite32(ext->Info.Address.Bus,
                          ext->Info.Address.Device,
                          ext->Info.Address.Function,
                          reg,
                          0xFFFFFFFFu);
            uint32_t maskLo = PciRawRead32(ext->Info.Address.Bus,
                                           ext->Info.Address.Device,
                                           ext->Info.Address.Function,
                                           reg);
            PciRawWrite32(ext->Info.Address.Bus,
                          ext->Info.Address.Device,
                          ext->Info.Address.Function,
                          reg,
                          original);

            bar->Type = PciBarMemory32;
            bar->Base = (uint64_t)(original & ~0xFu);
            bar->Length = PciBarSizeFromMask((uint64_t)maskLo, 0xFFFFFFF0ull);
        }
    }

    if (disabledCommand != savedCommand) {
        PciRawWrite16(ext->Info.Address.Bus,
                      ext->Info.Address.Device,
                      ext->Info.Address.Function,
                      PCI_REG_COMMAND_STATUS,
                      savedCommand);
    }
}

// Returns capability Offset.
static uint8_t PciFindCapabilityOffset(PCI_PDO_EXTENSION* ext) {
    uint16_t status = (uint16_t)(PciRawRead32(ext->Info.Address.Bus,
                                             ext->Info.Address.Device,
                                             ext->Info.Address.Function,
                                             PCI_REG_COMMAND_STATUS) >> 16);
    if ((status & PCI_STATUS_CAP_LIST) == 0)
        return 0;

    uint8_t cap = PciRawRead8(ext->Info.Address.Bus,
                              ext->Info.Address.Device,
                              ext->Info.Address.Function,
                              PCI_REG_CAP_PTR) & 0xFCu;
    return cap;
}

// Creates ids.
static void PciMakeIds(const PCI_DEVICE_INFO* info, IO_DEVICE_ID* ids, char storage[][IO_MAX_DEVICE_ID_CHARS], size_t* outCount) {
    size_t count = 0;

    snprintf(storage[count], IO_MAX_DEVICE_ID_CHARS,
             "PCI\\VEN_%04X&DEV_%04X&REV_%02X",
             info->VendorId, info->DeviceId, info->RevisionId);
    ids[count].BusType = IoBusPci;
    ids[count].Id = SysSvFromCstr(storage[count]);
    count++;

    snprintf(storage[count], IO_MAX_DEVICE_ID_CHARS,
             "PCI\\VEN_%04X&DEV_%04X",
             info->VendorId, info->DeviceId);
    ids[count].BusType = IoBusPci;
    ids[count].Id = SysSvFromCstr(storage[count]);
    count++;

    snprintf(storage[count], IO_MAX_DEVICE_ID_CHARS,
             "PCI\\CC_%02X%02X%02X",
             info->ClassCode, info->Subclass, info->ProgIf);
    ids[count].BusType = IoBusPci;
    ids[count].Id = SysSvFromCstr(storage[count]);
    count++;

    snprintf(storage[count], IO_MAX_DEVICE_ID_CHARS,
             "PCI\\CC_%02X%02X",
             info->ClassCode, info->Subclass);
    ids[count].BusType = IoBusPci;
    ids[count].Id = SysSvFromCstr(storage[count]);
    count++;

    snprintf(storage[count], IO_MAX_DEVICE_ID_CHARS,
             "PCI\\CC_%02X",
             info->ClassCode);
    ids[count].BusType = IoBusPci;
    ids[count].Id = SysSvFromCstr(storage[count]);
    count++;

    *outCount = count;
}

// Creates pdo.
static KSTATUS PciCreatePdo(uint8_t bus, uint8_t dev, uint8_t fn) {
    uint32_t id = PciRawRead32(bus, dev, fn, PCI_REG_VENDOR_DEVICE);
    uint16_t vendor = (uint16_t)(id & 0xFFFFu);
    if (vendor == PCI_VENDOR_INVALID)
        return STATUS_NOT_FOUND;

    uint32_t classReg = PciRawRead32(bus, dev, fn, PCI_REG_REVISION_CLASS);
    uint32_t hdrReg = PciRawRead32(bus, dev, fn, PCI_REG_HEADER_TYPE);
    uint32_t irqReg = PciRawRead32(bus, dev, fn, PCI_REG_INTERRUPT_LINE_PIN);
    uint32_t subsysReg = PciRawRead32(bus, dev, fn, PCI_REG_SUBSYSTEM);

    PCI_DEVICE_INFO info;
    memset(&info, 0, sizeof(info));
    info.Address.Segment = 0;
    info.Address.Bus = bus;
    info.Address.Device = dev;
    info.Address.Function = fn;
    info.VendorId = vendor;
    info.DeviceId = (uint16_t)(id >> 16);
    info.SubsystemVendorId = (uint16_t)(subsysReg & 0xFFFFu);
    info.SubsystemId = (uint16_t)(subsysReg >> 16);
    info.RevisionId = (uint8_t)(classReg & 0xFFu);
    info.ProgIf = (uint8_t)((classReg >> 8) & 0xFFu);
    info.Subclass = (uint8_t)((classReg >> 16) & 0xFFu);
    info.ClassCode = (uint8_t)((classReg >> 24) & 0xFFu);
    info.HeaderType = (uint8_t)((hdrReg >> 16) & 0xFFu);
    info.InterruptLine = (uint8_t)(irqReg & 0xFFu);
    info.InterruptPin = (uint8_t)((irqReg >> 8) & 0xFFu);

    IO_DEVICE_ID ids[5];
    char idStorage[5][IO_MAX_DEVICE_ID_CHARS];
    size_t idCount = 0;
    PciMakeIds(&info, ids, idStorage, &idCount);

    DEVICE_OBJECT* pdo = (DEVICE_OBJECT*)0;
    KSTATUS st = IoCreatePhysicalDevice(&gPciDriver,
                                        (DEVICE_OBJECT*)0,
                                        ids,
                                        idCount,
                                        sizeof(PCI_PDO_EXTENSION),
                                        &pdo);
    if (!KSUCCESS(st))
        return st;

    PCI_PDO_EXTENSION* ext = (PCI_PDO_EXTENSION*)pdo->DeviceExtension;
    memset(ext, 0, sizeof(*ext));
    ext->Info = info;
    PciDecodeBars(ext);
    ext->CapabilityOffset = PciFindCapabilityOffset(ext);
    pdo->u.Pdo.DeviceContext = ext;

    char name[IO_MAX_DEVICE_NAME_CHARS];
    snprintf(name, sizeof(name), "pci-%04X:%02X:%02X.%u",
             info.Address.Segment, info.Address.Bus, info.Address.Device, info.Address.Function);
    IoSetDeviceName(pdo, SysSvFromCstr(name));

    st = IoRegisterPhysicalDevice(pdo);
    if (!KSUCCESS(st)) {
        IoDeleteDevice(pdo);
        return st;
    }

    gPciDeviceCount++;
    DbgPrint("pci: %v ven=%04x dev=%04x class=%02x%02x%02x irq=%u\n",
             pdo->Name,
             info.VendorId,
             info.DeviceId,
             info.ClassCode,
             info.Subclass,
             info.ProgIf,
             info.InterruptLine);
    return STATUS_SUCCESS;
}

// Implements the enumerate Bus operation.
static void PciEnumerateBus(void) {
    for (uint32_t bus = 0; bus < 256u; bus++) {
        for (uint32_t dev = 0; dev < 32u; dev++) {
            uint32_t id0 = PciRawRead32((uint8_t)bus, (uint8_t)dev, 0, PCI_REG_VENDOR_DEVICE);
            if ((uint16_t)(id0 & 0xFFFFu) == PCI_VENDOR_INVALID)
                continue;

            uint8_t headerType = PciRawRead8((uint8_t)bus, (uint8_t)dev, 0, PCI_REG_HEADER_TYPE + 2u);
            uint8_t fnLimit = (headerType & PCI_HEADER_TYPE_MULTI) ? 8u : 1u;

            for (uint8_t fn = 0; fn < fnLimit; fn++) {
                KSTATUS st = PciCreatePdo((uint8_t)bus, (uint8_t)dev, fn);
                if (!KSUCCESS(st) && st != STATUS_NOT_FOUND) {
                    DbgPrint("pci: failed to create PDO for %02x:%02x.%u st=0x%x\n",
                             bus, dev, fn, st);
                }
            }
        }
    }
}

// Dispatches pnp.
static KSTATUS PciDispatchPnp(DEVICE_OBJECT* device, IRP* irp) {
    if (!device || !irp) return STATUS_INVALID_PARAMETER;
    IO_STACK_LOCATION* ios = IoGetCurrentIrpStackLocation(irp);
    if (!ios || ios->MajorFunction != IRP_MJ_PNP) return STATUS_INVALID_PARAMETER;

    switch ((IO_PNP_MINOR_FUNCTION)ios->MinorFunction) {
        case IRP_MN_START_DEVICE:
            device->PnpState = DevicePnpStarted;
            return STATUS_SUCCESS;
        case IRP_MN_QUERY_REMOVE_DEVICE:
        case IRP_MN_CANCEL_REMOVE_DEVICE:
        case IRP_MN_SURPRISE_REMOVAL:
        case IRP_MN_REMOVE_DEVICE:
            return STATUS_NOT_SUPPORTED;
        default:
            return STATUS_NOT_SUPPORTED;
    }
}

// Implements the driver Abi Version operation.
__declspec(dllexport) uint32_t PciDriverAbiVersion(void) {
    return 1u;
}

// Returns device Info.
__declspec(dllexport) KSTATUS PciGetDeviceInfo(DEVICE_OBJECT* pdo, PCI_DEVICE_INFO* outInfo) {
    PCI_PDO_EXTENSION* ext;
    if (!outInfo || !PciIsPdo(pdo, &ext)) return STATUS_INVALID_PARAMETER;
    *outInfo = ext->Info;
    return STATUS_SUCCESS;
}

// Returns config8.
__declspec(dllexport) KSTATUS PciReadConfig8(DEVICE_OBJECT* pdo, uint16_t offset, uint8_t* outValue) {
    PCI_PDO_EXTENSION* ext;
    if (!outValue || offset > 0xFFu || !PciIsPdo(pdo, &ext)) return STATUS_INVALID_PARAMETER;
    *outValue = PciRawRead8(ext->Info.Address.Bus, ext->Info.Address.Device, ext->Info.Address.Function, offset);
    return STATUS_SUCCESS;
}

// Returns config16.
__declspec(dllexport) KSTATUS PciReadConfig16(DEVICE_OBJECT* pdo, uint16_t offset, uint16_t* outValue) {
    PCI_PDO_EXTENSION* ext;
    if (!outValue || offset > 0xFEu || !PciIsPdo(pdo, &ext)) return STATUS_INVALID_PARAMETER;
    *outValue = PciRawRead16(ext->Info.Address.Bus, ext->Info.Address.Device, ext->Info.Address.Function, offset);
    return STATUS_SUCCESS;
}

// Returns config32.
__declspec(dllexport) KSTATUS PciReadConfig32(DEVICE_OBJECT* pdo, uint16_t offset, uint32_t* outValue) {
    PCI_PDO_EXTENSION* ext;
    if (!outValue || offset > 0xFCu || !PciIsPdo(pdo, &ext)) return STATUS_INVALID_PARAMETER;
    *outValue = PciRawRead32(ext->Info.Address.Bus, ext->Info.Address.Device, ext->Info.Address.Function, offset);
    return STATUS_SUCCESS;
}

// Updates config8.
__declspec(dllexport) KSTATUS PciWriteConfig8(DEVICE_OBJECT* pdo, uint16_t offset, uint8_t value) {
    PCI_PDO_EXTENSION* ext;
    if (offset > 0xFFu || !PciIsPdo(pdo, &ext)) return STATUS_INVALID_PARAMETER;
    PciRawWrite8(ext->Info.Address.Bus, ext->Info.Address.Device, ext->Info.Address.Function, offset, value);
    return STATUS_SUCCESS;
}

// Updates config16.
__declspec(dllexport) KSTATUS PciWriteConfig16(DEVICE_OBJECT* pdo, uint16_t offset, uint16_t value) {
    PCI_PDO_EXTENSION* ext;
    if (offset > 0xFEu || !PciIsPdo(pdo, &ext)) return STATUS_INVALID_PARAMETER;
    PciRawWrite16(ext->Info.Address.Bus, ext->Info.Address.Device, ext->Info.Address.Function, offset, value);
    return STATUS_SUCCESS;
}

// Updates config32.
__declspec(dllexport) KSTATUS PciWriteConfig32(DEVICE_OBJECT* pdo, uint16_t offset, uint32_t value) {
    PCI_PDO_EXTENSION* ext;
    if (offset > 0xFCu || !PciIsPdo(pdo, &ext)) return STATUS_INVALID_PARAMETER;
    PciRawWrite32(ext->Info.Address.Bus, ext->Info.Address.Device, ext->Info.Address.Function, offset, value);
    return STATUS_SUCCESS;
}

// Returns bar.
__declspec(dllexport) KSTATUS PciGetBar(DEVICE_OBJECT* pdo, uint8_t index, PCI_BAR* outBar) {
    PCI_PDO_EXTENSION* ext;
    if (!outBar || index >= PCI_MAX_BARS || !PciIsPdo(pdo, &ext)) return STATUS_INVALID_PARAMETER;
    if (ext->Bars[index].Type == PciBarUnused) return STATUS_NOT_FOUND;
    *outBar = ext->Bars[index];
    return STATUS_SUCCESS;
}

// Maps a PCI BAR through the private MMIO aperture with explicit cache attributes.
__declspec(dllexport) KSTATUS PciMapBar(DEVICE_OBJECT* pdo, uint8_t index, uint64_t pageFlags, void** outVirtual, uint64_t* outLength) {
    PCI_BAR bar;
    if (!outVirtual) return STATUS_INVALID_PARAMETER;
    *outVirtual = NULL;
    KSTATUS st = PciGetBar(pdo, index, &bar);
    if (!KSUCCESS(st)) return st;
    if (bar.Type != PciBarMemory32 && bar.Type != PciBarMemory64)
        return STATUS_NOT_SUPPORTED;
    if (!bar.Length)
        return STATUS_DEVICE_ERROR;
    if (bar.Base + bar.Length < bar.Base)
        return STATUS_INVALID_PARAMETER;

    uint64_t flags = pageFlags ? pageFlags : PciDefaultMmioFlags();
    flags |= PAGE_PRESENT;

    uint64_t physBase = bar.Base & ~(uint64_t)(PAGE_SIZE - 1u);
    uint64_t physOffset = bar.Base - physBase;
    uint64_t mapLength = PciAlignUp64(physOffset + bar.Length, PAGE_SIZE);

    uintptr_t virtBase = 0;
    st = PciAllocateMmioVa(mapLength, &virtBase);
    if (!KSUCCESS(st))
        return st;

    for (uint64_t off = 0; off < mapLength; off += PAGE_SIZE) {
        st = MmPteMapPage(virtBase + (uintptr_t)off, physBase + off, flags, MmPteMapFailIfPresent);
        if (!KSUCCESS(st))
            return st;
    }

    *outVirtual = (void*)(virtBase + (uintptr_t)physOffset);
    if (outLength) *outLength = bar.Length;
    return STATUS_SUCCESS;
}

// Enables device.
__declspec(dllexport) KSTATUS PciEnableDevice(DEVICE_OBJECT* pdo, uint32_t flags) {
    PCI_PDO_EXTENSION* ext;
    if (!PciIsPdo(pdo, &ext)) return STATUS_INVALID_PARAMETER;

    uint16_t cmd = PciRawRead16(ext->Info.Address.Bus,
                                ext->Info.Address.Device,
                                ext->Info.Address.Function,
                                PCI_REG_COMMAND_STATUS);
    if (flags & PCI_ENABLE_IO)        cmd |= PCI_COMMAND_IO_SPACE;
    if (flags & PCI_ENABLE_MEMORY)    cmd |= PCI_COMMAND_MEMORY_SPACE;
    if (flags & PCI_ENABLE_BUSMASTER) cmd |= PCI_COMMAND_BUS_MASTER;
    if (flags & PCI_ENABLE_INTX)      cmd &= (uint16_t)~PCI_COMMAND_INTX_DISABLE;
    if (flags & PCI_DISABLE_INTX)     cmd |= PCI_COMMAND_INTX_DISABLE;

    PciRawWrite16(ext->Info.Address.Bus,
                  ext->Info.Address.Device,
                  ext->Info.Address.Function,
                  PCI_REG_COMMAND_STATUS,
                  cmd);
    return STATUS_SUCCESS;
}

// Returns interrupt.
__declspec(dllexport) KSTATUS PciGetInterrupt(DEVICE_OBJECT* pdo, uint8_t* outIrq) {
    PCI_PDO_EXTENSION* ext;
    if (!outIrq || !PciIsPdo(pdo, &ext)) return STATUS_INVALID_PARAMETER;
    if (ext->Info.InterruptLine == 0xFFu) return STATUS_NOT_FOUND;
    *outIrq = ext->Info.InterruptLine;
    return STATUS_SUCCESS;
}

// Returns capability.
static uint8_t PciFindCapability(PCI_PDO_EXTENSION* ext, uint8_t capId) {
    if (!ext || !ext->CapabilityOffset)
        return 0;

    uint8_t cap = ext->CapabilityOffset & 0xFCu;
    for (uint32_t guard = 0; cap && guard < 48u; guard++) {
        uint32_t hdr = PciRawRead32(ext->Info.Address.Bus,
                                    ext->Info.Address.Device,
                                    ext->Info.Address.Function,
                                    cap);
        uint8_t id = (uint8_t)(hdr & 0xFFu);
        uint8_t next = (uint8_t)((hdr >> 8) & 0xFCu);
        if (id == capId)
            return cap;
        cap = next;
    }

    return 0;
}

// Enables msi.
__declspec(dllexport) KSTATUS PciEnableMsi(DEVICE_OBJECT* pdo, uint8_t irqSlot) {
    PCI_PDO_EXTENSION* ext;
    if (!PciIsPdo(pdo, &ext))
        return STATUS_INVALID_PARAMETER;
    if (irqSlot >= 24u)
        return STATUS_INVALID_PARAMETER;

    uint8_t cap = PciFindCapability(ext, PCI_CAP_ID_MSI);
    if (!cap)
        return STATUS_NOT_SUPPORTED;

    uint32_t bspLapic = HalGetBspLapicId();
    if (!HalApicUsingX2Apic() && bspLapic >= 256u)
        return STATUS_NOT_SUPPORTED;
    if (HalApicUsingX2Apic() && bspLapic >= 256u) {
        // Extended x2APIC MSI destinations need interrupt-remapping awareness.
        return STATUS_NOT_SUPPORTED;
    }

    uint8_t bus = ext->Info.Address.Bus;
    uint8_t dev = ext->Info.Address.Device;
    uint8_t fn = ext->Info.Address.Function;

    uint16_t control = PciRawRead16(bus, dev, fn, (uint16_t)(cap + 2u));
    bool is64 = (control & PCI_MSI_CTRL_64BIT) != 0;
    bool pvm = (control & PCI_MSI_CTRL_PVM) != 0;
    uint8_t vector = (uint8_t)(PCI_IRQ_VECTOR_BASE + irqSlot);
    uint32_t address = PCI_MSI_ADDRESS_BASE | ((bspLapic & 0xFFu) << 12);
    uint16_t data = vector;

    PciRawWrite32(bus, dev, fn, (uint16_t)(cap + 4u), address);
    if (is64) {
        PciRawWrite32(bus, dev, fn, (uint16_t)(cap + 8u), 0);
        PciRawWrite16(bus, dev, fn, (uint16_t)(cap + 12u), data);
        if (pvm)
            PciRawWrite32(bus, dev, fn, (uint16_t)(cap + 16u), 0);
    } else {
        PciRawWrite16(bus, dev, fn, (uint16_t)(cap + 8u), data);
        if (pvm)
            PciRawWrite32(bus, dev, fn, (uint16_t)(cap + 12u), 0);
    }

    control &= (uint16_t)~PCI_MSI_CTRL_MME_MASK;
    control |= PCI_MSI_CTRL_ENABLE;
    PciRawWrite16(bus, dev, fn, (uint16_t)(cap + 2u), control);

    uint16_t cmd = PciRawRead16(bus, dev, fn, PCI_REG_COMMAND_STATUS);
    cmd |= PCI_COMMAND_INTX_DISABLE;
    PciRawWrite16(bus, dev, fn, PCI_REG_COMMAND_STATUS, cmd);

    DbgPrint("pci: MSI enabled %02x:%02x.%u cap=0x%02x vector=0x%02x irqslot=%u dest=%u %s\n",
             bus, dev, fn, cap, vector, irqSlot, bspLapic, is64 ? "64-bit" : "32-bit");
    return STATUS_SUCCESS;
}

// Enables a single MSI vector, allocating the nt-lite IRQ slot internally.
__declspec(dllexport) KSTATUS PciEnableMsiAuto(DEVICE_OBJECT* pdo, uint8_t* outIrqSlot) {
    PCI_PDO_EXTENSION* ext;
    if (!outIrqSlot || !PciIsPdo(pdo, &ext))
        return STATUS_INVALID_PARAMETER;

    if (!PciFindCapability(ext, PCI_CAP_ID_MSI))
        return STATUS_NOT_SUPPORTED;

    uint8_t irq = 0xFFu;
    KSTATUS st = KeIrqAllocateMsiVector(&irq);
    if (!KSUCCESS(st))
        return st;

    st = PciEnableMsi(pdo, irq);
    if (!KSUCCESS(st))
        return st;

    *outIrqSlot = irq;
    return STATUS_SUCCESS;
}

// Implements the driver Entry operation.
KSTATUS DriverEntry(const HAL_API* hal) {
    if (!hal) return STATUS_INVALID_PARAMETER;
    gHal = hal;
    gPciDeviceCount = 0;

    IoInitializeDriverObject(&gPciDriver, SYS_SV_LITERAL("pci"));
    gPciDriver.MajorFunction[IRP_MJ_PNP] = PciDispatchPnp;

    KSTATUS st = IoRegisterDriverObject(&gPciDriver);
    if (!KSUCCESS(st)) return st;

    DbgPrint("pci: DriverEntry ABI=%u\n", PciDriverAbiVersion());
    PciEnumerateBus();
    DbgPrint("pci: enumeration complete, devices=%u\n", gPciDeviceCount);
    return STATUS_SUCCESS;
}