#include <ntlite/ddk/module.h>
// AHCI storage driver. This file owns controller attachment, drive discovery, command submission, and worker-driven IRP completion.

#include "ahci.h"

DRIVER_OBJECT  g_driver;

static uint8_t          g_nextDiskIndex;

#define AHCI_COMMAND_TIMEOUT_MS 10000ull
#define AHCI_PORT_ERROR_MASK \
    (PORT_IS_TFES | PORT_IS_HBFS | PORT_IS_HBDS | PORT_IS_IFS)
#define AHCI_PORT_TERMINAL_MASK (PORT_IS_DHRS | AHCI_PORT_ERROR_MASK)

static const IO_DEVICE_ID g_ahciIds[] = {
    { IoBusPci, { "PCI\\CC_010601", sizeof("PCI\\CC_010601") - 1 } },
};

static const IO_DRIVER_ID_TABLE g_ahciIdTable = {
    g_ahciIds,
    sizeof(g_ahciIds) / sizeof(g_ahciIds[0]),
};

// Helpers.

static KSTATUS AhciStartController(AHCI_CONTROLLER* ctl);

// Accessors for port structure memory at known CPU and device-address offsets.
static inline HBA_CMD_HEADER* cmd_list(AHCI_DRIVE* drv) {
    return (HBA_CMD_HEADER*)((uint8_t*)drv->structVirt + PORT_OFF_CMDLIST);
}

// Returns the command-list device address.
static inline uint64_t cmd_list_device_address(AHCI_DRIVE* drv) {
    return drv->structDeviceAddress + PORT_OFF_CMDLIST;
}

// Returns the received-FIS device address.
static inline uint64_t fis_device_address(AHCI_DRIVE* drv) {
    return drv->structDeviceAddress + PORT_OFF_FIS;
}

static inline HBA_CMD_TBL* cmd_tbl(AHCI_DRIVE* drv, int slot) {
    return (HBA_CMD_TBL*)((uint8_t*)drv->structVirt
                           + PORT_OFF_CMDTBLS
                           + (uint32_t)slot * PORT_CMDTBL_SIZE);
}

// Returns one command-table device address.
static inline uint64_t cmd_tbl_device_address(AHCI_DRIVE* drv, int slot) {
    return drv->structDeviceAddress + PORT_OFF_CMDTBLS
           + (uint64_t)slot * PORT_CMDTBL_SIZE;
}

// Port engine start / stop (polling safe).

// Port memory initialisation.

// Stops one drive worker before its queue and DMA storage are released.
static void AhciStopWorker(AHCI_DRIVE* drv) {
    if (!drv)
        return;

    drv->workerRunning = false;
    if (drv->irpQueue.Semaphore)
        KeSignalSemaphore(drv->irpQueue.Semaphore);
    if (drv->workerThread && drv->workerFinished)
        (void)KeWaitForEvent(drv->workerFinished, 0);

    if (drv->workerThread) {
        ObDereferenceObject((OBJECT*)drv->workerThread);
        drv->workerThread = (THREAD*)0;
    }
    if (drv->workerFinished) {
        ObDereferenceObject(&drv->workerFinished->Header);
        drv->workerFinished = (EVENT*)0;
    }
}

// Releases the non-DMA synchronization objects owned by one AHCI drive.
static void AhciReleaseDriveRuntime(AHCI_DRIVE* drv) {
    if (!drv)
        return;

    AhciStopWorker(drv);
    if (drv->removeLockReady) {
        IoDestroyRemoveLock(&drv->removeLock);
        drv->removeLockReady = false;
    }
    if (drv->irpQueue.Semaphore) {
        ObDereferenceObject(&drv->irpQueue.Semaphore->Header);
        drv->irpQueue.Semaphore = (SEMAPHORE*)0;
    }
    if (drv->completionEvent) {
        ObDereferenceObject(&drv->completionEvent->Header);
        drv->completionEvent = (EVENT*)0;
    }
}

// Releases the DMA buffers owned by one AHCI drive.
static void AhciReleaseDriveDma(AHCI_DRIVE* drv) {
    if (!drv)
        return;
    if (drv->port && drv->controller &&
        !drv->surpriseRemoved && !drv->controller->surpriseRemoved)
        port_stop(drv->port);
    IoDmaFreeCommonBuffer(drv->bounceDma);
    IoDmaFreeCommonBuffer(drv->structDma);
    drv->bounceDma = (IO_DMA_BUFFER*)0;
    drv->structDma = (IO_DMA_BUFFER*)0;
    drv->bounceDeviceAddress = 0;
    drv->structDeviceAddress = 0;
    drv->bounceVirt = (void*)0;
    drv->structVirt = (void*)0;
}

// Allocates and programs one port's command, FIS, and bounce buffers.
static KSTATUS port_init_memory(AHCI_DRIVE* drv) {
    if (!drv || !drv->controller || !drv->controller->dmaAdapter)
        return NTLSTATUS_INVALID_PARAMETER;

    KSTATUS status = IoDmaAllocateCommonBuffer(drv->controller->dmaAdapter,
                                                          PORT_PAGES_STRUCT * PAGE_SIZE,
                                                          PAGE_SIZE,
                                                          &drv->structDma);
    if (!KSUCCESS(status))
        return status;
    drv->structVirt = IoDmaGetCommonBufferAddress(drv->structDma);
    drv->structDeviceAddress = IoDmaGetCommonBufferDeviceAddress(drv->structDma);

    status = IoDmaAllocateCommonBuffer(drv->controller->dmaAdapter,
                                                   BOUNCE_BYTES,
                                                   PAGE_SIZE,
                                                   &drv->bounceDma);
    if (!KSUCCESS(status)) {
        AhciReleaseDriveDma(drv);
        return status;
    }
    drv->bounceVirt = IoDmaGetCommonBufferAddress(drv->bounceDma);
    drv->bounceDeviceAddress = IoDmaGetCommonBufferDeviceAddress(drv->bounceDma);

    HBA_PORT* port = drv->port;
    port_stop(port);

    port->clb  = (uint32_t)(cmd_list_device_address(drv) & 0xFFFFFFFF);
    port->clbu = (uint32_t)(cmd_list_device_address(drv) >> 32);
    port->fb   = (uint32_t)(fis_device_address(drv) & 0xFFFFFFFF);
    port->fbu  = (uint32_t)(fis_device_address(drv) >> 32);

    HBA_CMD_HEADER* hdrs = cmd_list(drv);
    for (int i = 0; i < NUM_CMD_SLOTS; i++) {
        hdrs[i].prdtl = PORT_PRDT_COUNT;
        uint64_t table = cmd_tbl_device_address(drv, i);
        hdrs[i].ctba  = (uint32_t)(table & 0xFFFFFFFF);
        hdrs[i].ctbau = (uint32_t)(table >> 32);
    }

    port->serr = (uint32_t)-1;
    port->is   = (uint32_t)-1;
    port_start(port);
    return NTLSTATUS_SUCCESS;
}

// Issue a command - polling fallback (used during early init / IDENTIFY only,
// before the scheduler and interrupt handler are live).
// Returns NTLSTATUS_SUCCESS, NTLSTATUS_DEVICE_ERROR, or NTLSTATUS_TIMEOUT.

static KSTATUS port_issue_cmd_poll(AHCI_DRIVE* drv) {
    HBA_PORT* port = drv->port;

    // Clear pending interrupts before issuing.
    port->is = (uint32_t)-1;

    // Wait for port to be idle (not BSY or DRQ).
    int spin = 0;
    while ((port->tfd & (ATA_SR_BUSY | ATA_SR_DRQ)) && spin < POLL_TIMEOUT)
        spin++;
    if (spin == POLL_TIMEOUT) {
        DbgPrint("AHCI: port hung (pre-issue)\n");
        return NTLSTATUS_TIMEOUT;
    }

    // Issue slot 0.
    port->ci = (1u << 0);

    // Poll until slot 0 clears, but never let early-init I/O hang forever.
    spin = 0;
    while ((port->ci & (1u << 0)) && spin < POLL_TIMEOUT) {
        if (port->is & AHCI_PORT_ERROR_MASK) {
            DbgPrint("AHCI: error IS=0x%x TFD=0x%x SERR=0x%x\n",
                     port->is, port->tfd, port->serr);
            port->serr = (uint32_t)-1;
            port->is   = (uint32_t)-1;
            return NTLSTATUS_DEVICE_ERROR;
        }
        spin++;
    }
    if (port->ci & (1u << 0)) {
        DbgPrint("AHCI: command timed out while polling CI=0x%x TFD=0x%x IS=0x%x SERR=0x%x\n",
                 port->ci, port->tfd, port->is, port->serr);
        return NTLSTATUS_TIMEOUT;
    }

    if (port->is & PORT_IS_TFES) {
        DbgPrint("AHCI: TFES after completion TFD=0x%x\n", port->tfd);
        port->is = (uint32_t)-1;
        return NTLSTATUS_DEVICE_ERROR;
    }

    return NTLSTATUS_SUCCESS;
}

// Recovers a port after a command that cannot be trusted to have retired.
static void AhciRecoverPort(AHCI_DRIVE* drv) {
    if (!drv || !drv->port)
        return;

    HBA_PORT* port = drv->port;
    uint32_t savedIe = port->ie;

    // Keep reset-generated FIS traffic from satisfying a command waiter.
    port->ie = 0;
    KSTATUS status = port_comreset(drv);
    port->serr = (uint32_t)-1;
    port->is = (uint32_t)-1;
    drv->lastIs = 0;
    if (drv->completionEvent)
        KeClearEvent(drv->completionEvent);
    port->ie = savedIe;

    if (!KSUCCESS(status)) {
        DbgPrint("AHCI: port recovery failed st=0x%x CI=0x%x TFD=0x%x SSTS=0x%x\n",
                 status, port->ci, port->tfd, port->ssts);
    }
}

// Interrupt bottom half. The ISR owns hardware acknowledgement and status
// capture; this DPC wakes the worker thread after the hard IRQ has unwound.

static void AhciCompletionDpc(KDPC* dpc, void* context,
                              void* systemArgument1, void* systemArgument2) {
    (void)dpc;
    (void)systemArgument1;
    (void)systemArgument2;

    AHCI_DRIVE* drv = (AHCI_DRIVE*)context;
    if (drv && drv->completionEvent)
        KeSignalEvent(drv->completionEvent);
}

// ISR - called from KeIrqDispatch on the AHCI PCI IRQ line.
// Walks all registered drives, checks if their port raised an interrupt,
// captures IS, clears it, and queues the per-drive completion DPC.

static void AhciIsr(void* context) {
    AHCI_CONTROLLER* ctl = (AHCI_CONTROLLER*)context;
    if (!ctl || !ctl->hba || ctl->removed) return;
    // GHC.IS tells us which ports are pending for this HBA.
    uint32_t pending = ctl->hba->is;
    if (!pending)
        return;

    for (uint32_t i = 0; i < ctl->driveCount; i++) {
        AHCI_DRIVE* drv = &ctl->drives[i];
        uint32_t portBit = (1u << drv->portNo);
        if (!(pending & portBit)) continue;

        // Capture and clear port IS before queueing the bottom half. The
        // waiter reads lastIs to check for errors without racing the next
        // command. OR in case hardware reports multiple causes before the
        // DPC releases the worker.
        uint32_t portIs = drv->port->is;
        drv->port->is   = portIs;          // W1C - write back to clear
        drv->lastIs    |= portIs;

        // Clear the HBA-level bit.
        ctl->hba->is = portBit;

        // DMA Setup, PIO Setup, and other progress interrupts are not
        // command completion. Only wake the worker for D2H completion or
        // an error that terminates the command.
        if (portIs & AHCI_PORT_TERMINAL_MASK)
            KeInsertQueueDpc(&drv->completionDpc, (void*)0, (void*)0);
    }

}


// Issue a command - interrupt-driven (runtime path).
// Caller has already built slot 0 via build_rw_cmd / port_identify.
// Blocks the current thread until the ISR signals completionEvent.

static KSTATUS port_issue_cmd(AHCI_DRIVE* drv) {
    if (drv->controller && drv->controller->pollingMode)
        return port_issue_cmd_poll(drv);

    HBA_PORT* port = drv->port;

    // A serialized slot must be idle before reuse. A stale active bit means a
    // prior command escaped normal completion and must be recovered first.
    if (port->ci & (1u << 0)) {
        DbgPrint("AHCI: slot 0 still active before issue CI=0x%x TFD=0x%x IS=0x%x SERR=0x%x\n",
                 port->ci, port->tfd, port->is, port->serr);
        AhciRecoverPort(drv);
        if (port->ci & (1u << 0))
            return NTLSTATUS_DEVICE_ERROR;
    }

    // Clear stale software and hardware completion state before submission.
    // The event can otherwise retain a late DPC signal from an older command.
    KeClearEvent(drv->completionEvent);
    port->is    = (uint32_t)-1;
    drv->lastIs = 0;

    // Wait for port idle.
    int spin = 0;
    while ((port->tfd & (ATA_SR_BUSY | ATA_SR_DRQ)) && spin < POLL_TIMEOUT)
        spin++;
    if (spin == POLL_TIMEOUT) {
        DbgPrint("AHCI: port hung before issue TFD=0x%x IS=0x%x SERR=0x%x\n",
                 port->tfd, port->is, port->serr);
        AhciRecoverPort(drv);
        return NTLSTATUS_TIMEOUT;
    }

    // Issue slot 0, then block with a finite deadline.
    port->ci = (1u << 0);
    EVENT* events[1] = { drv->completionEvent };
    KSTATUS waitStatus = KeWaitForMultipleEvents(events, 1, 0,
                                                  AHCI_COMMAND_TIMEOUT_MS,
                                                  (uint32_t*)0);

    // Merge any live status that was not consumed by the ISR. This lets a
    // missed IRQ still report the hardware error that accompanied it.
    uint32_t liveIs = port->is;
    uint32_t is = drv->lastIs | liveIs;
    if (liveIs)
        port->is = liveIs;

    if (is & AHCI_PORT_ERROR_MASK) {
        DbgPrint("AHCI: command error IS=0x%x CI=0x%x TFD=0x%x SERR=0x%x\n",
                 is, port->ci, port->tfd, port->serr);
        AhciRecoverPort(drv);
        return NTLSTATUS_DEVICE_ERROR;
    }

    if (waitStatus == NTLSTATUS_TIMEOUT) {
        if (!(port->ci & (1u << 0))) {
            DbgPrint("AHCI: completion IRQ missed; command retired IS=0x%x TFD=0x%x SERR=0x%x\n",
                     is, port->tfd, port->serr);
            return NTLSTATUS_SUCCESS;
        }

        DbgPrint("AHCI: command timeout CI=0x%x TFD=0x%x IS=0x%x SERR=0x%x\n",
                 port->ci, port->tfd, is, port->serr);
        AhciRecoverPort(drv);
        return NTLSTATUS_TIMEOUT;
    }

    if (waitStatus != NTLSTATUS_SUCCESS)
        return waitStatus;

    // An interrupt is only completion if hardware has also retired slot 0.
    if (port->ci & (1u << 0)) {
        DbgPrint("AHCI: premature completion IRQ CI=0x%x TFD=0x%x IS=0x%x SERR=0x%x\n",
                 port->ci, port->tfd, is, port->serr);
        AhciRecoverPort(drv);
        return NTLSTATUS_DEVICE_ERROR;
    }

    return NTLSTATUS_SUCCESS;
}

// Build a standard DMA read or write command in slot 0.

static void build_rw_cmd(AHCI_DRIVE* drv, uint64_t lba, uint32_t sectors,
                          uint64_t bufferDeviceAddress, bool write) {
    HBA_CMD_HEADER* hdr = cmd_list(drv);           // slot 0
    HBA_CMD_TBL* tbl = cmd_tbl(drv, 0);

    memset(tbl, 0, PORT_CMDTBL_SIZE);

    // How many PRDT entries we need (each covers up to 8K = 16 sectors).
    uint32_t prdtl = (sectors + 15) / 16;
    if (prdtl > PORT_PRDT_COUNT) prdtl = PORT_PRDT_COUNT; // capped (see caller)

    hdr->cfl   = sizeof(FIS_REG_H2D) / 4;
    hdr->w     = write ? 1 : 0;
    hdr->prdtl = (uint16_t)prdtl;
    hdr->prdbc = 0;

    // Fill PRDT entries.
    uint32_t remaining = sectors;
    uint64_t buf       = bufferDeviceAddress;
    for (uint32_t i = 0; i < prdtl; i++) {
        uint32_t sects = (remaining > 16) ? 16 : remaining;
        uint32_t bytes = sects * 512;

        tbl->prdt_entry[i].dba  = (uint32_t)(buf & 0xFFFFFFFF);
        tbl->prdt_entry[i].dbau = (uint32_t)(buf >> 32);
        tbl->prdt_entry[i].dbc  = bytes - 1;   // hardware adds 1
        tbl->prdt_entry[i].i    = (i == prdtl - 1) ? 1 : 0;

        buf       += bytes;
        remaining -= sects;
    }

    // Build H2D FIS.
    FIS_REG_H2D* fis = (FIS_REG_H2D*)tbl->cfis;
    memset(fis, 0, sizeof(*fis));
    fis->fis_type = FIS_TYPE_REG_H2D;
    fis->c        = 1;
    fis->command  = write ? ATA_CMD_WRITE_DMA_EX : ATA_CMD_READ_DMA_EX;
    fis->device   = 1u << 6;   // LBA mode

    fis->lba0 = (uint8_t)(lba);
    fis->lba1 = (uint8_t)(lba >> 8);
    fis->lba2 = (uint8_t)(lba >> 16);
    fis->lba3 = (uint8_t)(lba >> 24);
    fis->lba4 = (uint8_t)(lba >> 32);
    fis->lba5 = (uint8_t)(lba >> 40);

    fis->countl = (uint8_t)(sectors & 0xFF);
    fis->counth = (uint8_t)(sectors >> 8);
}

// Builds a non-data ATA FLUSH CACHE EXT command in slot 0.
static void build_flush_cmd(AHCI_DRIVE* drv) {
    HBA_CMD_HEADER* hdr = cmd_list(drv);
    HBA_CMD_TBL* tbl = cmd_tbl(drv, 0);

    memset(tbl, 0, PORT_CMDTBL_SIZE);
    hdr->cfl = sizeof(FIS_REG_H2D) / 4;
    hdr->w = 0;
    hdr->prdtl = 0;
    hdr->prdbc = 0;

    FIS_REG_H2D* fis = (FIS_REG_H2D*)tbl->cfis;
    memset(fis, 0, sizeof(*fis));
    fis->fis_type = FIS_TYPE_REG_H2D;
    fis->c = 1;
    fis->command = ATA_CMD_FLUSH_EXT;
}

// Flushes the drive write cache through the serialized command worker.
static KSTATUS do_flush(AHCI_DRIVE* drv) {
    if (!drv) return NTLSTATUS_INVALID_PARAMETER;
    build_flush_cmd(drv);
    return port_issue_cmd(drv);
}

// IDENTIFY command - fills a 512-byte buffer and returns it.

static KSTATUS port_identify(AHCI_DRIVE* drv) {
    // Reuse bounce buffer for IDENTIFY data (512 bytes).
    memset(drv->bounceVirt, 0, 512);

    HBA_CMD_HEADER* hdr = cmd_list(drv);
    HBA_CMD_TBL* tbl = cmd_tbl(drv, 0);

    memset(tbl, 0, PORT_CMDTBL_SIZE);

    hdr->cfl   = sizeof(FIS_REG_H2D) / 4;
    hdr->w     = 0;
    hdr->prdtl = 1;
    hdr->prdbc = 0;

    tbl->prdt_entry[0].dba  = (uint32_t)(drv->bounceDeviceAddress & 0xFFFFFFFF);
    tbl->prdt_entry[0].dbau = (uint32_t)(drv->bounceDeviceAddress >> 32);
    tbl->prdt_entry[0].dbc  = 512 - 1;
    tbl->prdt_entry[0].i    = 1;

    FIS_REG_H2D* fis = (FIS_REG_H2D*)tbl->cfis;
    memset(fis, 0, sizeof(*fis));
    fis->fis_type = FIS_TYPE_REG_H2D;
    fis->c        = 1;
    fis->command  = ATA_CMD_IDENTIFY;
    fis->device   = 0;

    // Poll: IDENTIFY runs before the scheduler and IRQ handler are live.
    KSTATUS st = port_issue_cmd_poll(drv);
    if (!KSUCCESS(st)) return st;

    // Parse sector count from IDENTIFY data.
    const uint16_t* id = (const uint16_t*)drv->bounceVirt;

    if (id[ATA_IDENT_W83] & (1u << 10)) {
        // 48-bit LBA supported.
        drv->sectorCount = (uint64_t)id[ATA_IDENT_W100]
                         | ((uint64_t)id[ATA_IDENT_W100 + 1] << 16)
                         | ((uint64_t)id[ATA_IDENT_W100 + 2] << 32)
                         | ((uint64_t)id[ATA_IDENT_W100 + 3] << 48);
    } else {
        // 28-bit LBA.
        drv->sectorCount = (uint64_t)id[ATA_IDENT_W60]
                         | ((uint64_t)id[ATA_IDENT_W60 + 1] << 16);
    }

    return NTLSTATUS_SUCCESS;
}

// Core read logic.
static KSTATUS do_read(AHCI_DRIVE* drv, uint64_t lba,
                        uint32_t count, void* buf) {
    uint8_t* out = (uint8_t*)buf;

    while (count > 0) {
        uint32_t batch = (count > BOUNCE_SECTORS) ? BOUNCE_SECTORS : count;

        build_rw_cmd(drv, lba, batch, drv->bounceDeviceAddress, false);
        KSTATUS st = port_issue_cmd(drv);
        if (!KSUCCESS(st)) return st;

        memcpy(out, drv->bounceVirt, (size_t)batch * 512);
        out   += (size_t)batch * 512;
        lba   += batch;
        count -= batch;
    }
    return NTLSTATUS_SUCCESS;
}

// Core write logic.
static KSTATUS do_write(AHCI_DRIVE* drv, uint64_t lba,
                         uint32_t count, const void* buf) {
    const uint8_t* in = (const uint8_t*)buf;

    while (count > 0) {
        uint32_t batch = (count > BOUNCE_SECTORS) ? BOUNCE_SECTORS : count;

        memcpy(drv->bounceVirt, in, (size_t)batch * 512);
        build_rw_cmd(drv, lba, batch, drv->bounceDeviceAddress, true);
        KSTATUS st = port_issue_cmd(drv);
        if (!KSUCCESS(st)) return st;

        in    += (size_t)batch * 512;
        lba   += batch;
        count -= batch;
    }
    return NTLSTATUS_SUCCESS;
}

// Async IRP dispatch.

static KSTATUS AhciGetIrpReadWrite(IRP* irp, bool* outWrite, uint64_t* outLba,
                                   uint32_t* outSectors, void** outBuffer) {
    if (!irp || !outWrite || !outLba || !outSectors || !outBuffer)
        return NTLSTATUS_INVALID_PARAMETER;
    IO_STACK_LOCATION* ios = IoGetCurrentIrpStackLocation(irp);
    if (!ios) return NTLSTATUS_INVALID_PARAMETER;
    if (ios->MajorFunction != IRP_MJ_READ && ios->MajorFunction != IRP_MJ_WRITE)
        return NTLSTATUS_INVALID_PARAMETER;
    if ((ios->Parameters.ReadWrite.ByteOffset & 511ull) ||
        (ios->Parameters.ReadWrite.Length & 511u))
        return NTLSTATUS_INVALID_PARAMETER;
    if (!irp->UserBuffer || ios->Parameters.ReadWrite.Length == 0)
        return NTLSTATUS_INVALID_PARAMETER;

    *outWrite = (ios->MajorFunction == IRP_MJ_WRITE);
    *outLba = ios->Parameters.ReadWrite.ByteOffset >> 9;
    *outSectors = ios->Parameters.ReadWrite.Length >> 9;
    *outBuffer = irp->UserBuffer;
    return NTLSTATUS_SUCCESS;
}

// AhciDispatchReadWrite - DEVICE_OBJECT dispatch entry for disk IRPs.
// Accepted IRPs are queued to the per-drive worker and completed later with
// IoCompleteRequest(). Rejected IRPs are returned synchronously and are not
// completed.
static KSTATUS AhciDispatchReadWrite(DEVICE_OBJECT* device, IRP* irp) {
    if (!device || !irp) return NTLSTATUS_INVALID_PARAMETER;
    AHCI_DRIVE* drv = (AHCI_DRIVE*)device->DeviceExtension;
    if (!drv) return NTLSTATUS_INVALID_PARAMETER;

    KSTATUS st = IoAcquireRemoveLock(&drv->removeLock, irp);
    if (!KSUCCESS(st))
        return st;

    bool write = false;
    uint64_t lba = 0;
    uint32_t sectors = 0;
    void* buffer = (void*)0;
    st = AhciGetIrpReadWrite(irp, &write, &lba, &sectors, &buffer);
    (void)write;
    (void)buffer;
    if (!KSUCCESS(st) || lba + sectors > drv->sectorCount) {
        IoReleaseRemoveLock(&drv->removeLock, irp);
        return KSUCCESS(st) ? NTLSTATUS_INVALID_PARAMETER : st;
    }

    irp->Status = NTLSTATUS_FAILURE;
    irp->Information = 0;
    IrpQueueEnqueue(&drv->irpQueue, irp);
    return NTLSTATUS_PENDING;
}

// Cancels release Remove Lock.
static void AhciCancelReleaseRemoveLock(IRP* irp, void* context) {
    AHCI_DRIVE* drv = (AHCI_DRIVE*)context;
    if (drv)
        IoReleaseRemoveLock(&drv->removeLock, irp);
}

// Dispatches drive Pnp.
static KSTATUS AhciDispatchDrivePnp(DEVICE_OBJECT* device, IRP* irp) {
    if (!device || !irp) return NTLSTATUS_INVALID_PARAMETER;
    AHCI_DRIVE* drv = (AHCI_DRIVE*)device->DeviceExtension;
    if (!drv) 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:
            device->PnpState = DevicePnpRemovePending;
            return NTLSTATUS_SUCCESS;
        case IRP_MN_CANCEL_REMOVE_DEVICE:
            device->PnpState = DevicePnpStarted;
            return NTLSTATUS_SUCCESS;
        case IRP_MN_SURPRISE_REMOVAL:
            drv->surpriseRemoved = true;
            device->PnpState = DevicePnpSurpriseRemoved;
            IrpQueueCancelAllWithCallback(&drv->irpQueue, NTLSTATUS_OBJECT_DEAD, AhciCancelReleaseRemoveLock, drv);
            return NTLSTATUS_SUCCESS;
        case IRP_MN_REMOVE_DEVICE:
            drv->removed = true;
            device->PnpState = DevicePnpDeleted;
            IrpQueueCancelAllWithCallback(&drv->irpQueue, NTLSTATUS_OBJECT_DEAD, AhciCancelReleaseRemoveLock, drv);
            if (drv->removeLockReady)
                IoReleaseRemoveLockAndWait(&drv->removeLock, irp);
            AhciReleaseDriveRuntime(drv);
            AhciReleaseDriveDma(drv);
            IoDeleteDevice(device);
            drv->deviceObject = (DEVICE_OBJECT*)0;
            return NTLSTATUS_SUCCESS;
        default:
            return NTLSTATUS_NOT_SUPPORTED;
    }
}

// Dispatches controller Pnp.
static KSTATUS AhciDispatchControllerPnp(DEVICE_OBJECT* device, IRP* irp) {
    if (!device || !irp) return NTLSTATUS_INVALID_PARAMETER;
    AHCI_CONTROLLER* ctl = (AHCI_CONTROLLER*)device->DeviceExtension;
    if (!ctl || ctl->magic != AHCI_CONTROLLER_MAGIC)
        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:
            if (ctl->started) {
                device->PnpState = DevicePnpStarted;
                return NTLSTATUS_SUCCESS;
            }
            return AhciStartController(ctl);

        case IRP_MN_QUERY_REMOVE_DEVICE:
            device->PnpState = DevicePnpRemovePending;
            return NTLSTATUS_SUCCESS;

        case IRP_MN_CANCEL_REMOVE_DEVICE:
            device->PnpState = DevicePnpStarted;
            return NTLSTATUS_SUCCESS;

        case IRP_MN_SURPRISE_REMOVAL:
            ctl->surpriseRemoved = true;
            device->PnpState = DevicePnpSurpriseRemoved;
            for (uint32_t i = 0; i < ctl->driveCount; i++) {
                AHCI_DRIVE* drv = &ctl->drives[i];
                drv->surpriseRemoved = true;
                if (drv->deviceObject)
                    (void)IoNotifyDiskSurpriseRemoved(drv->deviceObject->u.Disk.DiskIndex);
                else
                    IrpQueueCancelAllWithCallback(&drv->irpQueue, NTLSTATUS_OBJECT_DEAD, AhciCancelReleaseRemoveLock, drv);
            }
            return NTLSTATUS_SUCCESS;

        case IRP_MN_REMOVE_DEVICE:
            ctl->removed = true;
            device->PnpState = DevicePnpDeleted;
            for (uint32_t i = 0; i < ctl->driveCount; i++) {
                AHCI_DRIVE* drv = &ctl->drives[i];
                if (drv->deviceObject) {
                    (void)IoUnregisterDiskDevice(drv->deviceObject->u.Disk.DiskIndex);
                } else {
                    drv->removed = true;
                    IrpQueueCancelAllWithCallback(&drv->irpQueue, NTLSTATUS_OBJECT_DEAD, AhciCancelReleaseRemoveLock, drv);
                    AhciReleaseDriveRuntime(drv);
                    AhciReleaseDriveDma(drv);
                }
            }
            if (ctl->InterruptRequest.Interrupt)
                PciDisconnectInterrupts(ctl->pdo, &ctl->InterruptRequest, 1u);
            IoDmaPutAdapter(ctl->dmaAdapter);
            ctl->dmaAdapter = (IO_DMA_ADAPTER*)0;
            IoDetachDevice(device);
            IoDeleteDevice(device);
            return NTLSTATUS_SUCCESS;

        default:
            return NTLSTATUS_NOT_SUPPORTED;
    }
}

// Queues one write-cache flush on the per-drive worker.
static KSTATUS AhciDispatchFlush(DEVICE_OBJECT* device, IRP* irp) {
    if (!device || !irp) return NTLSTATUS_INVALID_PARAMETER;
    AHCI_DRIVE* drv = (AHCI_DRIVE*)device->DeviceExtension;
    if (!drv || drv->magic != AHCI_DRIVE_MAGIC)
        return NTLSTATUS_INVALID_PARAMETER;

    KSTATUS st = IoAcquireRemoveLock(&drv->removeLock, irp);
    if (!KSUCCESS(st))
        return st;

    IO_STACK_LOCATION* ios = IoGetCurrentIrpStackLocation(irp);
    if (!ios || ios->MajorFunction != IRP_MJ_FLUSH_BUFFERS) {
        IoReleaseRemoveLock(&drv->removeLock, irp);
        return NTLSTATUS_INVALID_PARAMETER;
    }

    irp->Status = NTLSTATUS_FAILURE;
    irp->Information = 0;
    IrpQueueEnqueue(&drv->irpQueue, irp);
    return NTLSTATUS_PENDING;
}

// Dispatches pnp.
static KSTATUS AhciDispatchPnp(DEVICE_OBJECT* device, IRP* irp) {
    if (!device || !device->DeviceExtension)
        return NTLSTATUS_INVALID_PARAMETER;

    uint32_t magic = *(uint32_t*)device->DeviceExtension;
    if (magic == AHCI_CONTROLLER_MAGIC)
        return AhciDispatchControllerPnp(device, irp);
    if (magic == AHCI_DRIVE_MAGIC)
        return AhciDispatchDrivePnp(device, irp);
    return NTLSTATUS_INVALID_PARAMETER;
}

// AhciWorker - kernel thread, one per drive.
// Dequeues IRPs and executes them one at a time using the existing
// single-slot interrupt-driven path. Completion fires the IRP's event
// and/or callback so the original caller can observe the result without
// ever having blocked in a driver call.

// NCQ remains serialized until the driver grows a multi-slot dispatcher:
//   1. Wait for a free slot and a pending IRP through the queue semaphore.
//   2. Pick the oldest IRP from the queue.
//   3. Grab a free bit in drv->slotMask (under the slot spinlock).
//   4. Copy write data into slotBounce[slot], build the NCQ FIS (FPDMA),
//      record slotIrp[slot] = irp, and issue via port->sact + port->ci.
//   5. Return immediately so the ISR can drive all completions.
static void AhciWorker(uintptr_t arg) {
    AHCI_DRIVE* drv = (AHCI_DRIVE*)arg;

    while (drv && drv->workerRunning) {
        // Block until AHCI dispatch signals a new IRP.
        KeWaitSemaphore(drv->irpQueue.Semaphore, KE_WAIT_INFINITE);
        if (!drv->workerRunning)
            break;

        IRP* irp = IrpQueueDequeue(&drv->irpQueue);
        if (!irp)
            continue;

        IO_STACK_LOCATION* ios = IoGetCurrentIrpStackLocation(irp);
        KSTATUS st = NTLSTATUS_INVALID_PARAMETER;
        size_t information = 0;

        if (ios && ios->MajorFunction == IRP_MJ_FLUSH_BUFFERS) {
            st = do_flush(drv);
        } else {
            bool write = false;
            uint64_t lba = 0;
            uint32_t sectors = 0;
            void* buffer = (void*)0;
            st = AhciGetIrpReadWrite(irp, &write, &lba, &sectors, &buffer);
            if (KSUCCESS(st)) {
                if (write)
                    st = do_write(drv, lba, sectors, buffer);
                else
                    st = do_read(drv, lba, sectors, buffer);
                if (KSUCCESS(st))
                    information = (size_t)sectors * 512u;
            }
        }

        // Publish result then fire completion hooks through the common I/O path.
        IoReleaseRemoveLock(&drv->removeLock, irp);
        IoCompleteRequest(irp, st, information);
    }

    if (drv && drv->workerFinished)
        KeSignalEvent(drv->workerFinished);
}

// AHCI controller initialisation.

// Before memory init: only check physical link. Signature is not yet valid.

// After COMRESET + FIS receive: check signature for SATA disk.

// Issue a COMRESET on the port, wait for link re-establishment and D2H FIS.
// FRE must already be set (port_init_memory does this) so the HBA can
// capture the signature FIS the device sends after the reset.

// Runtime commands are non-NCQ DMA commands. DHRS is their completion signal;
// setup/progress FIS interrupts must not wake the serialized command waiter.

// Starts controller.
static KSTATUS AhciStartController(AHCI_CONTROLLER* ctl) {
    if (!ctl || ctl->magic != AHCI_CONTROLLER_MAGIC || !ctl->pdo || !ctl->fdo)
        return NTLSTATUS_INVALID_PARAMETER;

    if (ctl->started)
        return NTLSTATUS_SUCCESS;

    KSTATUS st = PciGetDeviceInfo(ctl->pdo, &ctl->pciInfo);
    if (!KSUCCESS(st)) {
        DbgPrint("AHCI: PciGetDeviceInfo failed: 0x%x\n", st);
        return st;
    }

    DbgPrint("AHCI: controller at PCI %02x:%02x.%x vendor=%04x device=%04x\n",
             ctl->pciInfo.Address.Bus,
             ctl->pciInfo.Address.Device,
             ctl->pciInfo.Address.Function,
             ctl->pciInfo.VendorId,
             ctl->pciInfo.DeviceId);

    ctl->usingMsi = false;
    ctl->pollingMode = false;
    ctl->InterruptRequest.Routine = AhciIsr;
    ctl->InterruptRequest.Context = ctl;
    ctl->InterruptRequest.Affinity = 0;
    ctl->InterruptRequest.Interrupt = (KE_INTERRUPT*)0;

    st = PciGetBar(ctl->pdo, AHCI_PCI_BAR_ABAR, &ctl->abar);
    if (!KSUCCESS(st)) {
        DbgPrint("AHCI: BAR5/ABAR not present (0x%x)\n", st);
        return st;
    }
    if (ctl->abar.Type != PciBarMemory32 && ctl->abar.Type != PciBarMemory64) {
        DbgPrint("AHCI: BAR5/ABAR is not MMIO (type=%u)\n", ctl->abar.Type);
        return NTLSTATUS_DEVICE_ERROR;
    }

    st = PciMapBar(ctl->pdo,
                   AHCI_PCI_BAR_ABAR,
                   PAGE_PRESENT | PAGE_WRITE | PAGE_CACHE_UC,
                   (void**)&ctl->hba,
                   &ctl->abarLength);
    if (!KSUCCESS(st)) {
        DbgPrint("AHCI: PciMapBar(BAR5) failed: 0x%x\n", st);
        return st;
    }

    uint32_t enableFlags = PCI_ENABLE_MEMORY | PCI_ENABLE_BUSMASTER;
    enableFlags |= PCI_DISABLE_INTX;

    st = PciEnableDevice(ctl->pdo, enableFlags);
    if (!KSUCCESS(st)) {
        DbgPrint("AHCI: PciEnableDevice failed: 0x%x\n", st);
        return st;
    }

    DbgPrint("AHCI: ABAR phys=0x%016llX len=0x%llx version=0x%08x\n",
             (unsigned long long)ctl->abar.Base,
             (unsigned long long)ctl->abarLength,
             ctl->hba->vs);

    st = hba_reset(ctl->hba);
    if (!KSUCCESS(st))
        return st;

    uint32_t cap    = ctl->hba->cap;
    uint32_t nPorts = (cap & 0x1F) + 1;
    uint32_t nSlots = ((cap >> 8) & 0x1F) + 1;
    bool     s64a   = !!(cap & (1u << 31));
    DbgPrint("AHCI: cap=0x%08x  ports=%u  slots=%u  64bit=%d\n",
             cap, nPorts, nSlots, (int)s64a);

    IO_DMA_REQUIREMENTS dmaRequirements;
    memset(&dmaRequirements, 0, sizeof(dmaRequirements));
    dmaRequirements.MaximumDeviceAddress = s64a ? (uint64_t)-1 : 0xFFFFFFFFULL;
    dmaRequirements.Alignment = 2;
    dmaRequirements.MaximumSegmentLength = 1u << 22;
    dmaRequirements.MaximumSegments = PORT_PRDT_COUNT;
    st = IoDmaGetAdapter(ctl->fdo, &dmaRequirements, &ctl->dmaAdapter);
    if (!KSUCCESS(st)) {
        DbgPrint("AHCI: DMA adapter creation failed (0x%x)\n", st);
        return st;
    }

    // IoRegisterDiskDevice publishes arrival synchronously, and partition/filesystem
    // consumers may issue reads from that callback. Install the controller ISR before
    // any disk is published so an interrupt-driven drive can complete those reads.
    st = PciConnectInterrupts(ctl->pdo, &ctl->InterruptRequest, 1u, &ctl->usingMsi);
    if (!KSUCCESS(st)) {
        ctl->pollingMode = true;
        DbgPrint("AHCI: polling fallback status=0x%x\n", st);
    }

    uint32_t pi = ctl->hba->pi;
    for (uint32_t portno = 0; portno < 32 && ctl->driveCount < MAX_AHCI_DRIVES; portno++) {
        if (!(pi & (1u << portno))) continue;

        HBA_PORT* port = &ctl->hba->ports[portno];

        // Phase 1: physical link check.
        if (!port_link_up(port)) {
            DbgPrint("AHCI: port %u: no link (ssts=0x%08x)\n", portno, port->ssts);
            continue;
        }

        DbgPrint("AHCI: port %u: link up, initializing...\n", portno);

        uint32_t driveIdx = ctl->driveCount;
        AHCI_DRIVE* drv = &ctl->drives[driveIdx];
        memset(drv, 0, sizeof(*drv));
        drv->magic = AHCI_DRIVE_MAGIC;
        drv->controller = ctl;
        drv->port   = port;
        drv->portNo = (uint8_t)portno;

        // Phase 2: allocate command list, FIS buffer, bounce buffer.
        st = port_init_memory(drv);
        if (!KSUCCESS(st)) {
            DbgPrint("AHCI: port %u: memory init failed (0x%x)\n", portno, st);
            continue;
        }

        // Phase 3: COMRESET - forces device to re-send D2H FIS for valid SIG.
        st = port_comreset(drv);
        if (!KSUCCESS(st)) {
            DbgPrint("AHCI: port %u: COMRESET failed\n", portno);
            AhciReleaseDriveRuntime(drv);
            AhciReleaseDriveDma(drv);
            continue;
        }

        // Phase 4: filter to SATA disks.
        if (!port_is_sata_disk(port)) {
            DbgPrint("AHCI: port %u: not a SATA disk (sig=0x%08x), skipping\n",
                     portno, port->sig);
            AhciReleaseDriveRuntime(drv);
            AhciReleaseDriveDma(drv);
            continue;
        }

        // Phase 5: create per-drive completion event (auto-reset, initially clear).
        st = KeEventCreate(EventAutoReset, 0, &drv->completionEvent);
        if (!KSUCCESS(st)) {
            DbgPrint("AHCI: port %u: event create failed (0x%x)\n", portno, st);
            AhciReleaseDriveRuntime(drv);
            AhciReleaseDriveDma(drv);
            continue;
        }
        KeInitializeDpc(&drv->completionDpc, AhciCompletionDpc, drv);

        // Phase 5b: create the IRP queue semaphore and initialise the queue.
        // The semaphore starts at 0; each submitted IRP increments it by one.
        SEMAPHORE* irpSem = (SEMAPHORE*)0;
        st = KeSemaphoreCreate(0, 0x7FFFFFFF, &irpSem);
        if (!KSUCCESS(st)) {
            DbgPrint("AHCI: port %u: semaphore create failed (0x%x)\n", portno, st);
            AhciReleaseDriveRuntime(drv);
            AhciReleaseDriveDma(drv);
            continue;
        }
        IrpQueueInit(&drv->irpQueue, irpSem);
        st = IoInitializeRemoveLock(&drv->removeLock);
        if (!KSUCCESS(st)) {
            DbgPrint("AHCI: port %u: remove-lock init failed (0x%x)\n", portno, st);
            AhciReleaseDriveRuntime(drv);
            AhciReleaseDriveDma(drv);
            continue;
        }
        drv->removeLockReady = true;

        // Phase 6: IDENTIFY remains polled until the disk runtime path is ready.
        st = port_identify(drv);
        if (!KSUCCESS(st)) {
            DbgPrint("AHCI: port %u: IDENTIFY failed (0x%x)\n", portno, st);
            AhciReleaseDriveRuntime(drv);
            AhciReleaseDriveDma(drv);
            continue;
        }

        DbgPrint("AHCI: drive %d: %llu sectors (%llu MB)\n",
                 g_nextDiskIndex,
                 (unsigned long long)drv->sectorCount,
                 (unsigned long long)(drv->sectorCount / 2048));

        // Phase 7: create the disk DEVICE_OBJECT, but do not publish it until the
        // asynchronous runtime path is fully operational.
        st = IoCreateDevice(&g_driver, 0, &drv->deviceObject);
        if (!KSUCCESS(st)) {
            DbgPrint("AHCI: drive %d: IoCreateDevice failed (0x%x)\n", g_nextDiskIndex, st);
            AhciReleaseDriveRuntime(drv);
            AhciReleaseDriveDma(drv);
            continue;
        }
        drv->deviceObject->DeviceExtension = drv;
        drv->deviceObject->u.Disk.SectorCount = drv->sectorCount;
        drv->deviceObject->u.Disk.SectorSize = 512;
        drv->deviceObject->AlignmentRequirement = 512;

        // Phase 8: start the per-drive IRP worker before publishing the disk.
        // IoRegisterDiskDevice synchronously notifies partition consumers, which may
        // immediately submit reads that must have a live worker and completion path.
        st = KeEventCreate(EventManualReset, 0, &drv->workerFinished);
        if (KSUCCESS(st)) {
            drv->workerRunning = true;
            const KE_THREAD_CREATE_PARAMETERS workerParameters =
                KE_THREAD_CREATE_NAMED_NORMAL(PriorityHigh, "ahci.disk");
            st = KeCreateKernelThread((uintptr_t)AhciWorker,
                                      (uintptr_t)drv,
                                      &workerParameters,
                                      &drv->workerThread);
        }
        if (!KSUCCESS(st)) {
            drv->workerRunning = false;
            DbgPrint("AHCI: drive %u: worker creation failed (0x%x)\n", g_nextDiskIndex, st);
            IoDeleteDevice(drv->deviceObject);
            drv->deviceObject = (DEVICE_OBJECT*)0;
            AhciReleaseDriveRuntime(drv);
            AhciReleaseDriveDma(drv);
            continue;
        }

        // Make the initialized drive visible to the controller ISR, then arm its
        // interrupt source. No consumer can submit I/O until publication below.
        ctl->driveCount++;
        port->is = (uint32_t)-1;
        port->ie = AHCI_PORT_TERMINAL_MASK;

        // Phase 9: publish only after reads and writes can complete synchronously.
        uint8_t diskIndex = g_nextDiskIndex;
        st = IoRegisterDiskDevice(diskIndex, drv->deviceObject);
        if (KSUCCESS(st)) {
            DbgPrint("AHCI: drive %u registered\n", diskIndex);
            g_nextDiskIndex++;
        } else {
            DbgPrint("AHCI: IoRegisterDiskDevice(%u) failed (0x%x)\n", diskIndex, st);
            port->ie = 0;
            ctl->driveCount--;
            IoDeleteDevice(drv->deviceObject);
            drv->deviceObject = (DEVICE_OBJECT*)0;
            AhciReleaseDriveRuntime(drv);
            AhciReleaseDriveDma(drv);
            continue;
        }
    }

    if (ctl->driveCount == 0)
        DbgPrint("AHCI: no drives found on PCI %02x:%02x.%x\n",
                 ctl->pciInfo.Address.Bus,
                 ctl->pciInfo.Address.Device,
                 ctl->pciInfo.Address.Function);

    if (!ctl->pollingMode) {
        KE_INTERRUPT_INFO info;
        if (KSUCCESS(KeInterruptQuery(ctl->InterruptRequest.Interrupt, &info)))
            DbgPrint("AHCI: ISR cpu=%u vector=0x%02x mode=%s drives=%u\n",
                     info.CpuIndex, info.Vector, ctl->usingMsi ? "MSI/MSI-X" : "INTx",
                     (unsigned)ctl->driveCount);
    } else {
        DbgPrint("AHCI: controller started in polling mode; controller drives=%u\n",
                 (unsigned)ctl->driveCount);
    }
    ctl->started = true;
    ctl->fdo->PnpState = DevicePnpStarted;

    return NTLSTATUS_SUCCESS;
}

static KSTATUS AhciAddDevice(DRIVER_OBJECT* driver, DEVICE_OBJECT* physicalDeviceObject) {
    if (!driver || !physicalDeviceObject || physicalDeviceObject->Type != DeviceObjectPhysical)
        return NTLSTATUS_INVALID_PARAMETER;

    DEVICE_OBJECT* fdo = (DEVICE_OBJECT*)0;
    KSTATUS st = IoCreateDevice(driver, sizeof(AHCI_CONTROLLER), &fdo);
    if (!KSUCCESS(st))
        return st;

    AHCI_CONTROLLER* ctl = (AHCI_CONTROLLER*)fdo->DeviceExtension;
    memset(ctl, 0, sizeof(*ctl));
    ctl->magic = AHCI_CONTROLLER_MAGIC;
    ctl->fdo = fdo;
    ctl->pdo = physicalDeviceObject;
    ctl->pciInfo.InterruptLine = 0xFF;
    (void)IoSetDeviceName(fdo, SYS_SV_LITERAL("ahci-controller"));

    ctl->lowerDevice = IoAttachDeviceToDeviceStack(fdo, physicalDeviceObject);
    if (!ctl->lowerDevice) {
        IoDeleteDevice(fdo);
        return NTLSTATUS_FAILURE;
    }

    DbgPrint("AHCI: FDO attached to %v\n", physicalDeviceObject->Name);
    return NTLSTATUS_SUCCESS;
}

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

    DbgPrint("AHCI: driver loading...\n");

    g_nextDiskIndex = 0;

    memset(&g_driver, 0, sizeof(g_driver));
    IoInitializeDriverObject(&g_driver, SYS_SV_LITERAL("ahci"));
    g_driver.IdTable = &g_ahciIdTable;
    g_driver.DriverAddDevice = AhciAddDevice;
    g_driver.MajorFunction[IRP_MJ_READ] = AhciDispatchReadWrite;
    g_driver.MajorFunction[IRP_MJ_WRITE] = AhciDispatchReadWrite;
    g_driver.MajorFunction[IRP_MJ_FLUSH_BUFFERS] = AhciDispatchFlush;
    g_driver.MajorFunction[IRP_MJ_PNP] = AhciDispatchPnp;
    return NTLSTATUS_SUCCESS;
}

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