#include <ntlite/ddk/module.h>
// 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/ddk/ddk.h>
#include <ntlite/ddk/pci_caps.h>
#include <ntlite/ddi/io_driver.h>

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

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;
static SPINLOCK gPciMmioLock = SPINLOCK_INIT;

// 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 NTL_DDK_PAGE_PRESENT | NTL_DDK_PAGE_WRITE | NTL_DDK_PAGE_CACHE_UC | NTL_DDK_PAGE_NX;
}

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

    uint64_t alignedLength = PciAlignUp64(length, NTL_DDK_PAGE_SIZE);
    uintptr_t cookie = KeAcquireSpinlock(&gPciMmioLock);
    uintptr_t va = (uintptr_t)PciAlignUp64((uint64_t)gPciMmioNext, NTL_DDK_PAGE_SIZE);
    if (va < PCI_MMIO_VA_BASE || va >= PCI_MMIO_VA_LIMIT ||
        alignedLength > (uint64_t)(PCI_MMIO_VA_LIMIT - va)) {
        KeReleaseSpinlock(&gPciMmioLock, cookie);
        return NTLSTATUS_NO_MEMORY;
    }
    gPciMmioNext = va + (uintptr_t)alignedLength;
    KeReleaseSpinlock(&gPciMmioLock, cookie);
    *outVa = va;
    return NTLSTATUS_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 (offset > 0xFCu) return 0xFFFFFFFFu;
    return KePciReadConfig32(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 (offset > 0xFCu) return;
    KePciWriteConfig32(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 NTLSTATUS_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 NTLSTATUS_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 != NTLSTATUS_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 NTLSTATUS_INVALID_PARAMETER;
    IO_STACK_LOCATION* ios = IoGetCurrentIrpStackLocation(irp);
    if (!ios || ios->MajorFunction != IRP_MJ_PNP) return NTLSTATUS_INVALID_PARAMETER;

    switch ((IO_PNP_MINOR_FUNCTION)ios->MinorFunction) {
        case IRP_MN_START_DEVICE:
            device->PnpState = DevicePnpStarted;
            return NTLSTATUS_SUCCESS;
        case IRP_MN_QUERY_REMOVE_DEVICE:
        case IRP_MN_CANCEL_REMOVE_DEVICE:
        case IRP_MN_SURPRISE_REMOVAL:
        case IRP_MN_REMOVE_DEVICE:
            return NTLSTATUS_NOT_SUPPORTED;
        default:
            return NTLSTATUS_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 NTLSTATUS_INVALID_PARAMETER;
    *outInfo = ext->Info;
    return NTLSTATUS_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 NTLSTATUS_INVALID_PARAMETER;
    *outValue = PciRawRead8(ext->Info.Address.Bus, ext->Info.Address.Device, ext->Info.Address.Function, offset);
    return NTLSTATUS_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 NTLSTATUS_INVALID_PARAMETER;
    *outValue = PciRawRead16(ext->Info.Address.Bus, ext->Info.Address.Device, ext->Info.Address.Function, offset);
    return NTLSTATUS_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 NTLSTATUS_INVALID_PARAMETER;
    *outValue = PciRawRead32(ext->Info.Address.Bus, ext->Info.Address.Device, ext->Info.Address.Function, offset);
    return NTLSTATUS_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 NTLSTATUS_INVALID_PARAMETER;
    PciRawWrite8(ext->Info.Address.Bus, ext->Info.Address.Device, ext->Info.Address.Function, offset, value);
    return NTLSTATUS_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 NTLSTATUS_INVALID_PARAMETER;
    PciRawWrite16(ext->Info.Address.Bus, ext->Info.Address.Device, ext->Info.Address.Function, offset, value);
    return NTLSTATUS_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 NTLSTATUS_INVALID_PARAMETER;
    PciRawWrite32(ext->Info.Address.Bus, ext->Info.Address.Device, ext->Info.Address.Function, offset, value);
    return NTLSTATUS_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 NTLSTATUS_INVALID_PARAMETER;
    if (ext->Bars[index].Type == PciBarUnused) return NTLSTATUS_NOT_FOUND;
    *outBar = ext->Bars[index];
    return NTLSTATUS_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 NTLSTATUS_INVALID_PARAMETER;
    *outVirtual = NULL;
    KSTATUS st = PciGetBar(pdo, index, &bar);
    if (!KSUCCESS(st)) return st;
    if (bar.Type != PciBarMemory32 && bar.Type != PciBarMemory64)
        return NTLSTATUS_NOT_SUPPORTED;
    if (!bar.Length)
        return NTLSTATUS_DEVICE_ERROR;
    if (bar.Base + bar.Length < bar.Base)
        return NTLSTATUS_INVALID_PARAMETER;

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

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

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

    for (uint64_t off = 0; off < mapLength; off += NTL_DDK_PAGE_SIZE) {
        st = MmDdkMapPage(virtBase + (uintptr_t)off, physBase + off, flags, NtlDdkMapFailIfPresent);
        if (!KSUCCESS(st))
            return st;
    }

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

// Enables device.
__declspec(dllexport) KSTATUS PciEnableDevice(DEVICE_OBJECT* pdo, uint32_t flags) {
    PCI_PDO_EXTENSION* ext;
    if (!PciIsPdo(pdo, &ext)) return NTLSTATUS_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 NTLSTATUS_SUCCESS;
}

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

// Reads configuration through the PCI provider bound to this PDO.
static uint32_t PciCapabilityRead32(void* context, uint16_t offset) {
    const PCI_PDO_EXTENSION* ext = (const PCI_PDO_EXTENSION*)context;
    return PciRawRead32(ext->Info.Address.Bus,
                        ext->Info.Address.Device,
                        ext->Info.Address.Function,
                        offset);
}

// Finds one PCI capability using the shared DDK parser.
static uint8_t PciFindCapability(PCI_PDO_EXTENSION* ext, uint8_t capId) {
    if (!ext || !ext->CapabilityOffset)
        return 0;

    const NTL_PCI_CONFIG_READER reader = { ext, PciCapabilityRead32 };
    uint8_t offset = 0;
    if (!KSUCCESS(NtlPciFindCapability(&reader, ext->CapabilityOffset & 0xFCu,
                                       capId, &offset)))
        return 0;
    return offset;
}

// PCI owns message programming, while the kernel owns CPU/vector placement.
typedef struct PCI_MESSAGE_CONTEXT {
    DEVICE_OBJECT* Pdo;
    uint32_t TableOffset;
    uint16_t Capability;
    uint8_t Msix;
    uint8_t Reserved;
} PCI_MESSAGE_CONTEXT;

static KSTATUS PciUpdateMessageRoute(void* opaque,
    KE_INTERRUPT_MESSAGE_OPERATION operation, const NTL_DDK_MSI_MESSAGE* message) {
    PCI_MESSAGE_CONTEXT* ctx = (PCI_MESSAGE_CONTEXT*)opaque;
    PCI_PDO_EXTENSION* ext;
    if (!ctx || !PciIsPdo(ctx->Pdo, &ext)) return NTLSTATUS_INVALID_PARAMETER;
    uint8_t bus = ext->Info.Address.Bus, dev = ext->Info.Address.Device;
    uint8_t fn = ext->Info.Address.Function;
    if (ctx->Msix) {
        if (!ext->MsixTable || ctx->TableOffset > ext->MsixMappedLength ||
            ext->MsixMappedLength - ctx->TableOffset < 16u)
            return NTLSTATUS_DEVICE_ERROR;
        volatile uint32_t* entry = (volatile uint32_t*)(ext->MsixTable + ctx->TableOffset);
        if (operation == KeInterruptMessageQuiesce) {
            entry[3] = 1u;
        } else if (operation == KeInterruptMessageProgram && message) {
            entry[0] = (uint32_t)message->Address;
            entry[1] = (uint32_t)(message->Address >> 32);
            entry[2] = message->Data;
        } else if (operation == KeInterruptMessageResume) {
            entry[3] = 0u;
        } else return NTLSTATUS_INVALID_PARAMETER;
        (void)entry[3];
        return NTLSTATUS_SUCCESS;
    }
    uint16_t control = PciRawRead16(bus, dev, fn, ctx->Capability + 2u);
    bool is64 = (control & PCI_MSI_CTRL_64BIT) != 0;
    // MSI per-vector masking preserves pending messages during retargeting.
    // Devices without PVM retain their initial affinity until disconnect.
    if (!(control & PCI_MSI_CTRL_PVM)) return NTLSTATUS_NOT_SUPPORTED;
    uint16_t maskOffset = ctx->Capability + (is64 ? 16u : 12u);
    if (operation == KeInterruptMessageQuiesce) {
        uint32_t mask = PciRawRead32(bus, dev, fn, maskOffset);
        PciRawWrite32(bus, dev, fn, maskOffset, mask | 1u);
        (void)PciRawRead32(bus, dev, fn, maskOffset);
    } else if (operation == KeInterruptMessageProgram && message) {
        PciRawWrite32(bus, dev, fn, ctx->Capability + 4u, (uint32_t)message->Address);
        if (is64) {
            PciRawWrite32(bus, dev, fn, ctx->Capability + 8u,
                          (uint32_t)(message->Address >> 32));
            PciRawWrite16(bus, dev, fn, ctx->Capability + 12u, (uint16_t)message->Data);
        } else PciRawWrite16(bus, dev, fn, ctx->Capability + 8u, (uint16_t)message->Data);
    } else if (operation == KeInterruptMessageResume) {
        uint32_t mask = PciRawRead32(bus, dev, fn, maskOffset);
        PciRawWrite32(bus, dev, fn, maskOffset, mask & ~1u);
        (void)PciRawRead32(bus, dev, fn, maskOffset);
    } else return NTLSTATUS_INVALID_PARAMETER;
    (void)PciRawRead16(bus, dev, fn, ctx->Capability + 2u);
    return NTLSTATUS_SUCCESS;
}

static KSTATUS PciBindMessage(DEVICE_OBJECT* pdo, PCI_INTERRUPT_REQUEST* request,
                             uint8_t msix, uint16_t capability, uint32_t tableOffset) {
    PCI_MESSAGE_CONTEXT* ctx = (PCI_MESSAGE_CONTEXT*)MmAllocatePool(sizeof(*ctx));
    if (!ctx) return NTLSTATUS_NO_MEMORY;
    ctx->Pdo = pdo;
    ctx->Msix = msix;
    ctx->Capability = capability;
    ctx->TableOffset = tableOffset;
    ctx->Reserved = 0;
    KSTATUS st = KeInterruptBindMessageSource(request->Interrupt, PciUpdateMessageRoute, ctx);
    if (!KSUCCESS(st)) {
        MmFreePool(ctx);
        return st;
    }
    request->BusContext = ctx;
    return NTLSTATUS_SUCCESS;
}

// Programs a pre-allocated interrupt group. All routes exist before device delivery begins.
static KSTATUS PciConnectMsix(DEVICE_OBJECT* pdo, PCI_PDO_EXTENSION* ext,
                              PCI_INTERRUPT_REQUEST* requests, uint32_t count) {
    uint8_t cap = PciFindCapability(ext, PCI_CAP_ID_MSIX);
    if (!cap) return NTLSTATUS_NOT_SUPPORTED;
    uint8_t bus = ext->Info.Address.Bus, dev = ext->Info.Address.Device;
    uint8_t fn = ext->Info.Address.Function;
    uint16_t control = PciRawRead16(bus, dev, fn, cap + 2u);
    uint32_t size = (control & 0x7FFu) + 1u;
    if (count > size) return NTLSTATUS_NOT_SUPPORTED;
    uint32_t table = PciRawRead32(bus, dev, fn, cap + 4u);
    uint32_t bir = table & 7u, offset = table & ~7u;
    if (bir >= PCI_MAX_BARS) return NTLSTATUS_DEVICE_ERROR;
    PCI_BAR bar;
    KSTATUS st = PciGetBar(pdo, (uint8_t)bir, &bar);
    if (!KSUCCESS(st) || (bar.Type != PciBarMemory32 && bar.Type != PciBarMemory64) ||
        bar.Length < offset || count > (bar.Length - offset) / 16u)
        return NTLSTATUS_NOT_SUPPORTED;
    uint16_t command = PciRawRead16(bus, dev, fn, PCI_REG_COMMAND_STATUS);
    if (!(command & PCI_COMMAND_MEMORY_SPACE)) {
        DbgPrint("pci: MSI-X requires enabled memory decoding at %02x:%02x.%u\n",
                 bus, dev, fn);
        return NTLSTATUS_DEVICE_ERROR;
    }

    if (!ext->MsixTable) {
        volatile uint8_t* mapped = (volatile uint8_t*)0;
        uint64_t mappedLength = 0;
        st = PciMapBar(pdo, (uint8_t)bir,
                       NTL_DDK_PAGE_PRESENT | NTL_DDK_PAGE_WRITE | NTL_DDK_PAGE_CACHE_UC | NTL_DDK_PAGE_NX,
                       (void**)&mapped, &mappedLength);
        if (!KSUCCESS(st)) return st;
        ext->MsixTable = mapped;
        ext->MsixMappedLength = mappedLength;
        ext->MsixBarIndex = (uint8_t)bir;
    }
    // A persistent per-PDO mapping prevents repeated BAR allocations on reconnect.
    if (ext->MsixBarIndex != bir || ext->MsixMappedLength < offset ||
        count > (ext->MsixMappedLength - offset) / 16u)
        return NTLSTATUS_DEVICE_ERROR;
    volatile uint8_t* mapped = ext->MsixTable;

    // Mask the function and turn off stale firmware MSI before programming MSI-X.
    PciRawWrite16(bus, dev, fn, cap + 2u, (control | 0x4000u) & ~0x8000u);
    uint8_t msiCap = PciFindCapability(ext, PCI_CAP_ID_MSI);
    if (msiCap) {
        uint16_t msiControl = PciRawRead16(bus, dev, fn, msiCap + 2u);
        PciRawWrite16(bus, dev, fn, msiCap + 2u,
                      msiControl & (uint16_t)~PCI_MSI_CTRL_ENABLE);
    }
    for (uint32_t i = 0; i < count; i++) {
        st = KeInterruptReserveMessage(requests[i].Affinity, requests[i].Routine,
                                       requests[i].Context, &requests[i].Interrupt);
        if (!KSUCCESS(st)) goto failed;
        NTL_DDK_MSI_MESSAGE msg;
        st = KeInterruptComposeMessage(requests[i].Interrupt, &msg);
        if (!KSUCCESS(st)) goto failed;
        st = PciBindMessage(pdo, &requests[i], 1u, cap, offset + 16u * i);
        if (!KSUCCESS(st)) goto failed;
        volatile uint32_t* entry = (volatile uint32_t*)(mapped + offset + 16u * i);
        entry[3] = 1u;
        entry[0] = (uint32_t)msg.Address;
        entry[1] = (uint32_t)(msg.Address >> 32);
        entry[2] = msg.Data;
        entry[3] = 0u;
        // Read back the entry to drain posted MMIO writes before enabling MSI-X.
        (void)entry[3];
    }
    PciRawWrite16(bus, dev, fn, cap + 2u, (control | 0x8000u) & ~0x4000u);
    (void)PciRawRead16(bus, dev, fn, cap + 2u);
    PciRawWrite16(bus, dev, fn, PCI_REG_COMMAND_STATUS, command | PCI_COMMAND_INTX_DISABLE);
    for (uint32_t i = 0; i < count; i++) {
        st = KeInterruptActivateMessage(requests[i].Interrupt);
        if (!KSUCCESS(st)) goto failed;
    }
    return NTLSTATUS_SUCCESS;
failed:
    PciRawWrite16(bus, dev, fn, cap + 2u, control & ~0x8000u);
    for (uint32_t i = 0; i < count; i++) {
        if (requests[i].Interrupt) {
            KeInterruptDisconnect(requests[i].Interrupt);
            requests[i].Interrupt = (KE_INTERRUPT*)0;
        }
        if (requests[i].BusContext) {
            MmFreePool(requests[i].BusContext);
            requests[i].BusContext = 0;
        }
    }
    return st;
}

static KSTATUS PciConnectMsi(DEVICE_OBJECT* pdo, PCI_PDO_EXTENSION* ext,
                             PCI_INTERRUPT_REQUEST* request) {
    uint8_t cap = PciFindCapability(ext, PCI_CAP_ID_MSI);
    if (!cap) return NTLSTATUS_NOT_SUPPORTED;
    uint8_t bus = ext->Info.Address.Bus, dev = ext->Info.Address.Device;
    uint8_t fn = ext->Info.Address.Function;
    KSTATUS st = KeInterruptReserveMessage(request->Affinity, request->Routine,
                                           request->Context, &request->Interrupt);
    if (!KSUCCESS(st)) return st;
    NTL_DDK_MSI_MESSAGE msg;
    st = KeInterruptComposeMessage(request->Interrupt, &msg);
    if (!KSUCCESS(st)) goto failed;
    st = PciBindMessage(pdo, request, 0u, cap, 0);
    if (!KSUCCESS(st)) goto failed;
    uint8_t msixCap = PciFindCapability(ext, PCI_CAP_ID_MSIX);
    if (msixCap) {
        uint16_t msixControl = PciRawRead16(bus, dev, fn, msixCap + 2u);
        PciRawWrite16(bus, dev, fn, msixCap + 2u,
                      (msixControl | 0x4000u) & (uint16_t)~0x8000u);
    }
    uint16_t control = PciRawRead16(bus, dev, fn, cap + 2u);
    bool is64 = (control & PCI_MSI_CTRL_64BIT) != 0;
    bool pvm = (control & PCI_MSI_CTRL_PVM) != 0;
    PciRawWrite32(bus, dev, fn, cap + 4u, (uint32_t)msg.Address);
    if (is64) {
        PciRawWrite32(bus, dev, fn, cap + 8u, (uint32_t)(msg.Address >> 32));
        PciRawWrite16(bus, dev, fn, cap + 12u, (uint16_t)msg.Data);
        if (pvm) PciRawWrite32(bus, dev, fn, cap + 16u, 0);
    } else {
        PciRawWrite16(bus, dev, fn, cap + 8u, (uint16_t)msg.Data);
        if (pvm) PciRawWrite32(bus, dev, fn, cap + 12u, 0);
    }
    control = (control & ~(uint16_t)PCI_MSI_CTRL_MME_MASK) | PCI_MSI_CTRL_ENABLE;
    PciRawWrite16(bus, dev, fn, cap + 2u, control);
    uint16_t command = PciRawRead16(bus, dev, fn, PCI_REG_COMMAND_STATUS);
    PciRawWrite16(bus, dev, fn, PCI_REG_COMMAND_STATUS, command | PCI_COMMAND_INTX_DISABLE);
    st = KeInterruptActivateMessage(request->Interrupt);
    if (!KSUCCESS(st)) goto failed;
    return NTLSTATUS_SUCCESS;
failed:
    {
        uint16_t previous = PciRawRead16(bus, dev, fn, cap + 2u);
        PciRawWrite16(bus, dev, fn, cap + 2u,
                      previous & (uint16_t)~PCI_MSI_CTRL_ENABLE);
    }
    KeInterruptDisconnect(request->Interrupt);
    request->Interrupt = (KE_INTERRUPT*)0;
    if (request->BusContext) {
        MmFreePool(request->BusContext);
        request->BusContext = 0;
    }
    return st;
}

__declspec(dllexport) KSTATUS PciConnectInterrupts(DEVICE_OBJECT* pdo,
    PCI_INTERRUPT_REQUEST* requests, uint32_t count, bool* usingMessages) {
    PCI_PDO_EXTENSION* ext;
    if (!PciIsPdo(pdo, &ext) || !requests || !count || !usingMessages)
        return NTLSTATUS_INVALID_PARAMETER;
    *usingMessages = false;
    for (uint32_t i = 0; i < count; i++) {
        if (!requests[i].Routine) return NTLSTATUS_INVALID_PARAMETER;
        requests[i].Interrupt = (KE_INTERRUPT*)0;
        requests[i].BusContext = 0;
    }
    KSTATUS st = PciConnectMsix(pdo, ext, requests, count);
    if (KSUCCESS(st)) { *usingMessages = true; return st; }
    if (count > 1u) return st;
    st = PciConnectMsi(pdo, ext, requests);
    if (KSUCCESS(st)) { *usingMessages = true; return st; }
    uint8_t line = ext->Info.InterruptLine;
    uint8_t pin = ext->Info.InterruptPin;
    if (pin < 1u || pin > 4u || line == 0u || line == 0xFFu) {
        DbgPrint("pci: INTx resource unavailable for %02x:%02x.%u pin=%u line=%u\n",
                 ext->Info.Address.Bus, ext->Info.Address.Device,
                 ext->Info.Address.Function, pin, line);
        return NTLSTATUS_NOT_FOUND;
    }
    // The Interrupt Line register is a firmware hint until PCI _PRT routing is available.
    DbgPrint("pci: INTx legacy route %02x:%02x.%u pin=%u line=%u\n",
             ext->Info.Address.Bus, ext->Info.Address.Device,
             ext->Info.Address.Function, pin, line);
    if (line < 16u)
        st = KeInterruptConnectIsa(line, 0x0Fu, requests[0].Affinity,
            requests[0].Routine, requests[0].Context, &requests[0].Interrupt);
    else
        st = KeInterruptConnectGsi(line, 0x0Fu, requests[0].Affinity,
            requests[0].Routine, requests[0].Context, &requests[0].Interrupt);
    if (KSUCCESS(st)) {
        uint16_t command = PciRawRead16(ext->Info.Address.Bus,
            ext->Info.Address.Device, ext->Info.Address.Function, PCI_REG_COMMAND_STATUS);
        PciRawWrite16(ext->Info.Address.Bus, ext->Info.Address.Device,
            ext->Info.Address.Function, PCI_REG_COMMAND_STATUS,
            command & (uint16_t)~PCI_COMMAND_INTX_DISABLE);
    }
    return st;
}

__declspec(dllexport) void PciDisconnectInterrupts(DEVICE_OBJECT* pdo,
    PCI_INTERRUPT_REQUEST* requests, uint32_t count) {
    PCI_PDO_EXTENSION* ext;
    if (!PciIsPdo(pdo, &ext) || !requests) return;
    // Close every route before disabling PCI delivery so an affinity update
    // cannot resume the device after its disconnect has begun.
    for (uint32_t i = 0; i < count; i++)
        if (requests[i].Interrupt) KeInterruptBeginDisconnect(requests[i].Interrupt);
    uint8_t bus = ext->Info.Address.Bus, dev = ext->Info.Address.Device;
    uint8_t fn = ext->Info.Address.Function;
    uint8_t msi = PciFindCapability(ext, PCI_CAP_ID_MSI);
    uint8_t msix = PciFindCapability(ext, PCI_CAP_ID_MSIX);
    if (msi) {
        uint16_t control = PciRawRead16(bus, dev, fn, msi + 2u);
        PciRawWrite16(bus, dev, fn, msi + 2u, control & ~PCI_MSI_CTRL_ENABLE);
    }
    if (msix) {
        uint16_t control = PciRawRead16(bus, dev, fn, msix + 2u);
        PciRawWrite16(bus, dev, fn, msix + 2u, (control | 0x4000u) & ~0x8000u);
    }
    uint16_t command = PciRawRead16(bus, dev, fn, PCI_REG_COMMAND_STATUS);
    PciRawWrite16(bus, dev, fn, PCI_REG_COMMAND_STATUS, command | PCI_COMMAND_INTX_DISABLE);
    for (uint32_t i = 0; i < count; i++) {
        if (requests[i].Interrupt) {
            KeInterruptDisconnect(requests[i].Interrupt);
            requests[i].Interrupt = (KE_INTERRUPT*)0;
        }
        if (requests[i].BusContext) {
            MmFreePool(requests[i].BusContext);
            requests[i].BusContext = 0;
        }
    }
}

static KSTATUS DriverInitialize(NTL_MODULE* module) {
    (void)module;

    gPciDeviceCount = 0;

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

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

static DRIVER_OBJECT* const gModuleIoDrivers[] = { &gPciDriver };
NTL_DECLARE_IO_MODULE(DriverInitialize, gModuleIoDrivers);