#include <ntlite/ddk/module.h>
// ATA compatibility storage driver. This file owns legacy IDE probing and PIO disk access for older virtual hardware.

#include "ata.h"

static DRIVER_OBJECT gAtaDriver;
static ATA_CONTROLLER* gAtaControllers[ATA_MAX_CONTROLLERS];
static uint32_t gAtaControllerCount;

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

static const IO_DRIVER_ID_TABLE gAtaIdTable = {
    gAtaIds,
    sizeof(gAtaIds) / sizeof(gAtaIds[0]),
};

static uint8_t AtaIn8(ATA_CHANNEL* ch, uint8_t reg) {
    return __inbyte((unsigned short)(ch->CommandBase + reg));
}

static uint16_t AtaIn16(ATA_CHANNEL* ch, uint8_t reg) {
    return __inword((unsigned short)(ch->CommandBase + reg));
}

static void AtaOut8(ATA_CHANNEL* ch, uint8_t reg, uint8_t value) {
    __outbyte((unsigned short)(ch->CommandBase + reg), value);
}

static void AtaOut16(ATA_CHANNEL* ch, uint8_t reg, uint16_t value) {
    __outword((unsigned short)(ch->CommandBase + reg), value);
}

static uint8_t AtaAltStatus(ATA_CHANNEL* ch) {
    return __inbyte((unsigned short)(ch->ControlBase + ATA_REG_ALT_STATUS));
}

static void AtaDeviceControl(ATA_CHANNEL* ch, uint8_t value) {
    __outbyte((unsigned short)(ch->ControlBase + ATA_REG_DEVICE_CONTROL), value);
}

static void AtaDelay400ns(ATA_CHANNEL* ch) {
    (void)AtaAltStatus(ch);
    (void)AtaAltStatus(ch);
    (void)AtaAltStatus(ch);
    (void)AtaAltStatus(ch);
}

// Returns the monotonic DDK timer in microseconds.
static uint64_t AtaQueryTicks(void) {
    return KeQueryTicks();
}

// Delays briefly between ATA status samples without involving the scheduler.
static void AtaPollDelay(void) {
    KeStallExecution(ATA_POLL_INTERVAL_US);
}

// Returns whether a status byte represents an unimplemented/floating legacy channel.
static bool AtaStatusIsAbsent(uint8_t status) {
    return status == 0x00u || status == 0xFFu;
}

// Waits until one microsecond deadline expires while polling alternate status.
static KSTATUS AtaWaitStatus(ATA_CHANNEL* ch, uint64_t timeoutUs, bool requireDrq, bool checkErrors,
                             uint8_t* finalStatusOut) {
    uint64_t start = AtaQueryTicks();
    uint64_t deadline = start > UINT64_MAX - timeoutUs ? UINT64_MAX : start + timeoutUs;
    uint64_t pollBudget = timeoutUs / ATA_POLL_INTERVAL_US + 2u;

    for (;;) {
        uint8_t status = AtaAltStatus(ch);
        if (finalStatusOut) *finalStatusOut = status;
        if (AtaStatusIsAbsent(status)) return NTLSTATUS_NOT_FOUND;
        if (!(status & ATA_STATUS_BSY)) {
            if (checkErrors && (status & (ATA_STATUS_ERR | ATA_STATUS_DF)))
                return NTLSTATUS_DEVICE_ERROR;
            if (!requireDrq || (status & ATA_STATUS_DRQ))
                return NTLSTATUS_SUCCESS;
        }
        if (AtaQueryTicks() >= deadline || pollBudget-- == 0u)
            return NTLSTATUS_TIMEOUT;
        AtaPollDelay();
    }
}

// Waits boundedly for BSY to clear using elapsed hardware time rather than scheduler iterations.
static KSTATUS AtaWaitNotBusy(ATA_CHANNEL* ch) {
    return AtaWaitStatus(ch, ATA_POLL_TIMEOUT_US, false, true, (uint8_t*)0);
}

// Waits boundedly for a data phase while rejecting absent and failed devices early.
static KSTATUS AtaWaitDrq(ATA_CHANNEL* ch) {
    return AtaWaitStatus(ch, ATA_POLL_TIMEOUT_US, true, true, (uint8_t*)0);
}

// Emits the task-file state needed to distinguish selection, signature, and command failures.
static void AtaLogProbeState(const ATA_DEVICE* dev, const char* stage, KSTATUS result) {
    if (!dev || !dev->Channel) return;
    ATA_CHANNEL* ch = dev->Channel;
    DbgPrint("ATA: channel %u %s identify %s result=0x%x status=0x%02x error=0x%02x lba1=0x%02x lba2=0x%02x\n",
             ch->Index,
             dev->Drive ? "slave" : "master",
             stage ? stage : "stage",
             result,
             AtaAltStatus(ch),
             AtaIn8(ch, ATA_REG_ERROR),
             AtaIn8(ch, ATA_REG_LBA1),
             AtaIn8(ch, ATA_REG_LBA2));
}

// Resets one legacy ATA channel before probing so firmware state cannot leak into IDENTIFY.
static KSTATUS AtaResetChannel(ATA_CHANNEL* ch) {
    if (!ch) return NTLSTATUS_INVALID_PARAMETER;

    uint8_t before = AtaAltStatus(ch);
    DbgPrint("ATA: channel %u reset begin status=0x%02x\n", ch->Index, before);

    AtaDeviceControl(ch, ATA_DEVICE_CONTROL_NIEN | ATA_DEVICE_CONTROL_SRST);
    KeStallExecution(ATA_RESET_ASSERT_US);
    AtaDeviceControl(ch, ATA_DEVICE_CONTROL_NIEN);
    KeStallExecution(ATA_RESET_SETTLE_US);

    uint8_t finalStatus = AtaAltStatus(ch);
    if (AtaStatusIsAbsent(finalStatus)) {
        DbgPrint("ATA: channel %u reset found no responding device status=0x%02x\n",
                 ch->Index, finalStatus);
        return NTLSTATUS_NOT_FOUND;
    }

    KSTATUS status = AtaWaitStatus(ch, ATA_RESET_TIMEOUT_US, false, false, &finalStatus);
    DbgPrint("ATA: channel %u reset complete result=0x%x status=0x%02x\n",
             ch->Index, status, finalStatus);
    return status;
}

static void AtaSelectDrive(ATA_DEVICE* dev) {
    ATA_CHANNEL* ch = dev->Channel;
    AtaOut8(ch, ATA_REG_DRIVE_HEAD, (uint8_t)(0xA0u | (dev->Drive ? 0x10u : 0u)));
    AtaDelay400ns(ch);
}

// Implements IDENTIFY with stage diagnostics and finite hardware deadlines.
static KSTATUS AtaIdentify(ATA_DEVICE* dev) {
    ATA_CHANNEL* ch = dev->Channel;
    uint16_t ident[256];
    memset(ident, 0, sizeof(ident));

    KSTATUS st = KeWaitSemaphore(ch->Lock, KE_WAIT_INFINITE);
    if (!KSUCCESS(st))
        return st;
    DbgPrint("ATA: channel %u %s identify lock acquired\n",
             ch->Index, dev->Drive ? "slave" : "master");

    AtaSelectDrive(dev);
    DbgPrint("ATA: channel %u %s identify selected status=0x%02x\n",
             ch->Index, dev->Drive ? "slave" : "master", AtaAltStatus(ch));

    st = AtaWaitStatus(ch, ATA_POLL_TIMEOUT_US, false, false, (uint8_t*)0);
    if (!KSUCCESS(st)) {
        AtaLogProbeState(dev, "pre-command wait", st);
        KeSignalSemaphore(ch->Lock);
        return st;
    }

    AtaOut8(ch, ATA_REG_SECTOR_COUNT, 0);
    AtaOut8(ch, ATA_REG_LBA0, 0);
    AtaOut8(ch, ATA_REG_LBA1, 0);
    AtaOut8(ch, ATA_REG_LBA2, 0);
    AtaOut8(ch, ATA_REG_COMMAND, ATA_CMD_IDENTIFY_DEVICE);
    AtaDelay400ns(ch);

    uint8_t status = AtaAltStatus(ch);
    DbgPrint("ATA: channel %u %s identify command issued status=0x%02x\n",
             ch->Index, dev->Drive ? "slave" : "master", status);
    if (AtaStatusIsAbsent(status)) {
        AtaLogProbeState(dev, "absent after command", NTLSTATUS_NOT_FOUND);
        KeSignalSemaphore(ch->Lock);
        return NTLSTATUS_NOT_FOUND;
    }

    st = AtaWaitNotBusy(ch);
    if (!KSUCCESS(st)) {
        AtaLogProbeState(dev, "busy timeout/error", st);
        KeSignalSemaphore(ch->Lock);
        return st;
    }

    uint8_t lba1 = AtaIn8(ch, ATA_REG_LBA1);
    uint8_t lba2 = AtaIn8(ch, ATA_REG_LBA2);
    if (lba1 != 0 || lba2 != 0) {
        AtaLogProbeState(dev, "non-ATA signature", NTLSTATUS_NOT_SUPPORTED);
        KeSignalSemaphore(ch->Lock);
        return NTLSTATUS_NOT_SUPPORTED;
    }

    st = AtaWaitDrq(ch);
    if (!KSUCCESS(st)) {
        AtaLogProbeState(dev, "DRQ timeout/error", st);
        KeSignalSemaphore(ch->Lock);
        return st;
    }

    for (uint32_t i = 0; i < 256; i++)
        ident[i] = AtaIn16(ch, ATA_REG_DATA);

    AtaLogProbeState(dev, "data complete", NTLSTATUS_SUCCESS);
    KeSignalSemaphore(ch->Lock);
    DbgPrint("ATA: channel %u %s identify lock released capabilities=0x%04x command_sets=0x%04x\n",
             ch->Index, dev->Drive ? "slave" : "master",
             ident[ATA_IDENT_CAPABILITIES], ident[ATA_IDENT_COMMAND_SETS]);

    if (!(ident[ATA_IDENT_CAPABILITIES] & ATA_IDENT_LBA_SUPPORTED))
        return NTLSTATUS_NOT_SUPPORTED;

    dev->SupportsLba48 = (ident[ATA_IDENT_COMMAND_SETS] & ATA_IDENT_LBA48_SUPPORTED) != 0;
    if (dev->SupportsLba48) {
        uint64_t sectors = 0;
        sectors |= (uint64_t)ident[100];
        sectors |= (uint64_t)ident[101] << 16;
        sectors |= (uint64_t)ident[102] << 32;
        sectors |= (uint64_t)ident[103] << 48;
        dev->SectorCount = sectors;
    } else {
        dev->SectorCount = ((uint32_t)ident[ATA_IDENT_LBA28_SECTORS_HI] << 16) |
                           ident[ATA_IDENT_LBA28_SECTORS_LO];
    }

    DbgPrint("ATA: channel %u %s identify parsed lba48=%u sectors=%llu\n",
             ch->Index, dev->Drive ? "slave" : "master",
             dev->SupportsLba48 ? 1u : 0u, (unsigned long long)dev->SectorCount);
    return dev->SectorCount ? NTLSTATUS_SUCCESS : NTLSTATUS_NOT_FOUND;
}

static KSTATUS AtaProgramLba(ATA_DEVICE* dev, uint64_t lba, uint16_t count, bool write) {
    ATA_CHANNEL* ch = dev->Channel;

    if (dev->SupportsLba48) {
        AtaOut8(ch, ATA_REG_DRIVE_HEAD, (uint8_t)(0x40u | (dev->Drive ? 0x10u : 0u)));
        AtaOut8(ch, ATA_REG_SECTOR_COUNT, (uint8_t)(count >> 8));
        AtaOut8(ch, ATA_REG_LBA0, (uint8_t)(lba >> 24));
        AtaOut8(ch, ATA_REG_LBA1, (uint8_t)(lba >> 32));
        AtaOut8(ch, ATA_REG_LBA2, (uint8_t)(lba >> 40));
        AtaOut8(ch, ATA_REG_SECTOR_COUNT, (uint8_t)count);
        AtaOut8(ch, ATA_REG_LBA0, (uint8_t)lba);
        AtaOut8(ch, ATA_REG_LBA1, (uint8_t)(lba >> 8));
        AtaOut8(ch, ATA_REG_LBA2, (uint8_t)(lba >> 16));
        AtaOut8(ch, ATA_REG_COMMAND, write ? ATA_CMD_WRITE_SECTORS_EXT : ATA_CMD_READ_SECTORS_EXT);
        return NTLSTATUS_SUCCESS;
    }

    if (lba > 0x0FFFFFFFULL || count > 256)
        return NTLSTATUS_INVALID_PARAMETER;

    AtaOut8(ch, ATA_REG_DRIVE_HEAD, (uint8_t)(0xE0u | (dev->Drive ? 0x10u : 0u) | ((lba >> 24) & 0x0Fu)));
    AtaOut8(ch, ATA_REG_SECTOR_COUNT, (uint8_t)count);
    AtaOut8(ch, ATA_REG_LBA0, (uint8_t)lba);
    AtaOut8(ch, ATA_REG_LBA1, (uint8_t)(lba >> 8));
    AtaOut8(ch, ATA_REG_LBA2, (uint8_t)(lba >> 16));
    AtaOut8(ch, ATA_REG_COMMAND, write ? ATA_CMD_WRITE_SECTORS : ATA_CMD_READ_SECTORS);
    return NTLSTATUS_SUCCESS;
}

// Returns pio.
static KSTATUS AtaReadPio(ATA_DEVICE* dev, uint64_t lba, uint32_t count, void* buffer) {
    if (!dev || !buffer)
        return NTLSTATUS_INVALID_PARAMETER;
    if (count == 0)
        return NTLSTATUS_SUCCESS;
    if (lba + count > dev->SectorCount)
        return NTLSTATUS_INVALID_PARAMETER;

    ATA_CHANNEL* ch = dev->Channel;
    uint16_t* out = (uint16_t*)buffer;

    KSTATUS st = KeWaitSemaphore(ch->Lock, KE_WAIT_INFINITE);
    if (!KSUCCESS(st))
        return st;

    while (count > 0) {
        uint16_t batch = (count > 256u) ? 256u : (uint16_t)count;
        uint16_t encodedCount = (!dev->SupportsLba48 && batch == 256u) ? 0u : batch;

        AtaSelectDrive(dev);
        st = AtaWaitNotBusy(ch);
        if (!KSUCCESS(st)) break;

        st = AtaProgramLba(dev, lba, encodedCount, false);
        if (!KSUCCESS(st)) break;

        for (uint32_t s = 0; s < batch; s++) {
            st = AtaWaitDrq(ch);
            if (!KSUCCESS(st)) break;
            for (uint32_t w = 0; w < ATA_SECTOR_SIZE / 2u; w++)
                *out++ = AtaIn16(ch, ATA_REG_DATA);
        }
        if (!KSUCCESS(st)) break;

        lba += batch;
        count -= batch;
    }

    KeSignalSemaphore(ch->Lock);
    return st;
}

// Flushes the requested operation.
static KSTATUS AtaFlush(ATA_DEVICE* dev) {
    ATA_CHANNEL* ch = dev->Channel;
    KSTATUS st = AtaWaitNotBusy(ch);
    if (!KSUCCESS(st))
        return st;
    AtaOut8(ch, ATA_REG_COMMAND, dev->SupportsLba48 ? ATA_CMD_CACHE_FLUSH_EXT : ATA_CMD_CACHE_FLUSH);
    return AtaWaitNotBusy(ch);
}

// Flushes one device while holding the shared channel lock.
static KSTATUS AtaFlushSerialized(ATA_DEVICE* dev) {
    if (!dev || !dev->Channel)
        return NTLSTATUS_INVALID_PARAMETER;

    ATA_CHANNEL* ch = dev->Channel;
    KSTATUS st = KeWaitSemaphore(ch->Lock, KE_WAIT_INFINITE);
    if (!KSUCCESS(st))
        return st;

    AtaSelectDrive(dev);
    st = AtaFlush(dev);
    KeSignalSemaphore(ch->Lock);
    return st;
}

// Updates pio.
static KSTATUS AtaWritePio(ATA_DEVICE* dev, uint64_t lba, uint32_t count, const void* buffer) {
    if (!dev || !buffer)
        return NTLSTATUS_INVALID_PARAMETER;
    if (count == 0)
        return NTLSTATUS_SUCCESS;
    if (lba + count > dev->SectorCount)
        return NTLSTATUS_INVALID_PARAMETER;

    ATA_CHANNEL* ch = dev->Channel;
    const uint16_t* in = (const uint16_t*)buffer;

    KSTATUS st = KeWaitSemaphore(ch->Lock, KE_WAIT_INFINITE);
    if (!KSUCCESS(st))
        return st;

    while (count > 0) {
        uint16_t batch = (count > 256u) ? 256u : (uint16_t)count;
        uint16_t encodedCount = (!dev->SupportsLba48 && batch == 256u) ? 0u : batch;

        AtaSelectDrive(dev);
        st = AtaWaitNotBusy(ch);
        if (!KSUCCESS(st)) break;

        st = AtaProgramLba(dev, lba, encodedCount, true);
        if (!KSUCCESS(st)) break;

        for (uint32_t s = 0; s < batch; s++) {
            st = AtaWaitDrq(ch);
            if (!KSUCCESS(st)) break;
            for (uint32_t w = 0; w < ATA_SECTOR_SIZE / 2u; w++)
                AtaOut16(ch, ATA_REG_DATA, *in++);
        }
        if (!KSUCCESS(st)) break;

        st = AtaFlush(dev);
        if (!KSUCCESS(st)) break;

        lba += batch;
        count -= batch;
    }

    KeSignalSemaphore(ch->Lock);
    return st;
}

static KSTATUS AtaGetIrpReadWrite(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;
}

static void AtaWorker(uintptr_t arg) {
    ATA_DEVICE* dev = (ATA_DEVICE*)arg;
    if (!dev)
        return;

    DbgPrint("ATA: channel %u %s worker online\n",
             dev->Channel ? dev->Channel->Index : 0u, dev->Drive ? "slave" : "master");
    while (dev->WorkerRunning) {
        KSTATUS waitStatus = KeWaitSemaphore(dev->IrpQueue.Semaphore, KE_WAIT_INFINITE);
        if (!KSUCCESS(waitStatus))
            continue;
        if (!dev->WorkerRunning)
            break;

        IRP* irp = IrpQueueDequeue(&dev->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 = AtaFlushSerialized(dev);
        } else {
            bool write = false;
            uint64_t lba = 0;
            uint32_t sectors = 0;
            void* buffer = (void*)0;
            st = AtaGetIrpReadWrite(irp, &write, &lba, &sectors, &buffer);
            if (KSUCCESS(st)) {
                if (lba + sectors > dev->SectorCount) {
                    st = NTLSTATUS_INVALID_PARAMETER;
                } else {
                    st = write ? AtaWritePio(dev, lba, sectors, buffer)
                               : AtaReadPio(dev, lba, sectors, buffer);
                    if (KSUCCESS(st))
                        information = (size_t)sectors * ATA_SECTOR_SIZE;
                }
            }
        }

        IoReleaseRemoveLock(&dev->RemoveLock, irp);
        IoCompleteRequest(irp, st, information);
    }

    DbgPrint("ATA: channel %u %s worker stopped\n",
             dev->Channel ? dev->Channel->Index : 0u, dev->Drive ? "slave" : "master");
    if (dev->WorkerFinished)
        KeSignalEvent(dev->WorkerFinished);
}

// Stops one ATA disk worker and releases the worker-owned synchronization references.
static void AtaStopWorker(ATA_DEVICE* dev) {
    if (!dev)
        return;

    dev->WorkerRunning = false;
    if (dev->IrpQueue.Semaphore)
        KeSignalSemaphore(dev->IrpQueue.Semaphore);
    if (dev->WorkerThread && dev->WorkerFinished)
        (void)KeWaitForEvent(dev->WorkerFinished, 0);

    if (dev->WorkerThread) {
        ObDereferenceObject((OBJECT*)dev->WorkerThread);
        dev->WorkerThread = (THREAD*)0;
    }
    if (dev->WorkerFinished) {
        ObDereferenceObject(&dev->WorkerFinished->Header);
        dev->WorkerFinished = (EVENT*)0;
    }
    if (dev->IrpQueue.Semaphore) {
        ObDereferenceObject(&dev->IrpQueue.Semaphore->Header);
        dev->IrpQueue.Semaphore = (SEMAPHORE*)0;
    }
}

// Starts one ATA disk worker before the disk is published to synchronous discovery consumers.
static KSTATUS AtaStartWorker(ATA_DEVICE* dev) {
    if (!dev || !dev->IrpQueue.Semaphore)
        return NTLSTATUS_INVALID_PARAMETER;

    KSTATUS status = KeEventCreate(EventManualReset, 0, &dev->WorkerFinished);
    if (!KSUCCESS(status))
        return status;

    dev->WorkerRunning = true;
    const KE_THREAD_CREATE_PARAMETERS workerParameters =
        KE_THREAD_CREATE_NAMED_NORMAL(PriorityHigh, "ata.disk");
    status = KeCreateKernelThread((uintptr_t)AtaWorker, (uintptr_t)dev,
                                  &workerParameters, &dev->WorkerThread);
    if (!KSUCCESS(status)) {
        dev->WorkerRunning = false;
        AtaStopWorker(dev);
        return status;
    }
    return NTLSTATUS_SUCCESS;
}

// Dispatches read Write.
KSTATUS AtaDispatchReadWrite(DEVICE_OBJECT* device, IRP* irp) {
    if (!device || !irp || !device->DeviceExtension)
        return NTLSTATUS_INVALID_PARAMETER;

    ATA_DEVICE* dev = (ATA_DEVICE*)device->DeviceExtension;
    if (dev->Magic != ATA_DEVICE_MAGIC || !dev->Present)
        return NTLSTATUS_INVALID_PARAMETER;

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

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

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

// Queues one cache flush on the device worker.
static KSTATUS AtaDispatchFlush(DEVICE_OBJECT* device, IRP* irp) {
    if (!device || !irp || !device->DeviceExtension)
        return NTLSTATUS_INVALID_PARAMETER;

    ATA_DEVICE* dev = (ATA_DEVICE*)device->DeviceExtension;
    if (dev->Magic != ATA_DEVICE_MAGIC || !dev->Present)
        return NTLSTATUS_INVALID_PARAMETER;

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

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

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

// Returns free Disk Index.
static uint8_t AtaFindFreeDiskIndex(void) {
    for (uint8_t i = 0; i < 8; i++) {
        if (!IoDiskDevicePresent(i))
            return i;
    }
    return 0xFFu;
}

static uint16_t AtaBarIoBase(ATA_CONTROLLER* ctl, uint8_t barIndex, uint16_t fallback) {
    PCI_BAR bar;
    KSTATUS st = PciGetBar(ctl->Pdo, barIndex, &bar);
    if (KSUCCESS(st) && bar.Type == PciBarIo && bar.Base && bar.Base <= 0xFFFFu)
        return (uint16_t)bar.Base;
    return fallback;
}

// Initializes channel.
static KSTATUS AtaInitChannel(ATA_CONTROLLER* ctl, uint8_t index) {
    ATA_CHANNEL* ch = &ctl->Channels[index];
    memset(ch, 0, sizeof(*ch));
    ch->Controller = ctl;
    ch->Index = index;

    if (index == 0) {
        ch->CommandBase = AtaBarIoBase(ctl, 0, ATA_PRIMARY_LEGACY_CMD_BASE);
        ch->ControlBase = AtaBarIoBase(ctl, 1, ATA_PRIMARY_LEGACY_CTL_BASE);
        ch->Irq = ATA_PRIMARY_LEGACY_IRQ;
    } else {
        ch->CommandBase = AtaBarIoBase(ctl, 2, ATA_SECONDARY_LEGACY_CMD_BASE);
        ch->ControlBase = AtaBarIoBase(ctl, 3, ATA_SECONDARY_LEGACY_CTL_BASE);
        ch->Irq = ATA_SECONDARY_LEGACY_IRQ;
    }

    if (ctl->PciInfo.ProgIf & (index == 0 ? 0x01u : 0x04u)) {
        uint8_t irq = 0xFFu;
        if (KSUCCESS(PciGetInterrupt(ctl->Pdo, &irq)) && irq != 0xFFu)
            ch->Irq = irq;
        ch->NativeMode = true;
    }

    KSTATUS st = KeSemaphoreCreate(1, 1, &ch->Lock);
    if (!KSUCCESS(st))
        return st;

    AtaDeviceControl(ch, ATA_DEVICE_CONTROL_NIEN); // nIEN: polled I/O for now.
    DbgPrint("ATA: channel %u cmd=0x%04x ctl=0x%04x irq=%u mode=%s\n",
             index,
             ch->CommandBase,
             ch->ControlBase,
             ch->Irq,
             ch->NativeMode ? "native" : "compat");
    return NTLSTATUS_SUCCESS;
}

// Registers disk.
static KSTATUS AtaRegisterDisk(ATA_DEVICE* dev) {
    DEVICE_OBJECT* disk = (DEVICE_OBJECT*)0;
    KSTATUS st = IoCreateDevice(&gAtaDriver, 0, &disk);
    if (!KSUCCESS(st))
        return st;

    disk->DeviceExtension = dev;
    disk->u.Disk.SectorCount = dev->SectorCount;
    disk->u.Disk.SectorSize = ATA_SECTOR_SIZE;
    disk->AlignmentRequirement = ATA_SECTOR_SIZE;
    (void)IoSetDeviceName(disk, SYS_SV_LITERAL("ata-disk"));

    uint8_t index = AtaFindFreeDiskIndex();
    if (index == 0xFFu) {
        IoDeleteDevice(disk);
        return NTLSTATUS_NO_MEMORY;
    }

    st = IoRegisterDiskDevice(index, disk);
    if (!KSUCCESS(st)) {
        IoDeleteDevice(disk);
        return st;
    }

    dev->DeviceObject = disk;
    dev->DiskIndex = index;
    DbgPrint("ATA: drive %u registered channel=%u %s sectors=%llu (%llu MB)\n",
             index,
             dev->Channel->Index,
             dev->Drive ? "slave" : "master",
             (unsigned long long)dev->SectorCount,
             (unsigned long long)(dev->SectorCount / 2048));
    return NTLSTATUS_SUCCESS;
}

// Probes drive.
static KSTATUS AtaProbeDrive(ATA_CONTROLLER* ctl, ATA_CHANNEL* ch, uint8_t drive) {
    ATA_DEVICE* dev = &ch->Devices[drive];
    memset(dev, 0, sizeof(*dev));
    dev->Magic = ATA_DEVICE_MAGIC;
    dev->Controller = ctl;
    dev->Channel = ch;
    dev->Drive = drive;
    dev->DiskIndex = 0xFFu;

    KSTATUS st = AtaIdentify(dev);
    if (!KSUCCESS(st))
        return st;

    st = IoInitializeRemoveLock(&dev->RemoveLock);
    if (!KSUCCESS(st))
        return st;

    SEMAPHORE* sem = (SEMAPHORE*)0;
    st = KeSemaphoreCreate(0, 0x7FFFFFFF, &sem);
    if (!KSUCCESS(st))
        return st;
    IrpQueueInit(&dev->IrpQueue, sem);

    DbgPrint("ATA: channel %u %s worker start before disk publication\n",
             ch->Index, drive ? "slave" : "master");
    st = AtaStartWorker(dev);
    if (!KSUCCESS(st)) {
        DbgPrint("ATA: channel %u %s worker start failed (0x%x)\n",
                 ch->Index, drive ? "slave" : "master", st);
        AtaStopWorker(dev);
        IoDestroyRemoveLock(&dev->RemoveLock);
        return st;
    }

    dev->Present = true;
    DbgPrint("ATA: channel %u %s publishing disk sectors=%llu lba48=%u\n",
             ch->Index, drive ? "slave" : "master",
             (unsigned long long)dev->SectorCount, dev->SupportsLba48 ? 1u : 0u);
    st = AtaRegisterDisk(dev);
    if (!KSUCCESS(st)) {
        DbgPrint("ATA: channel %u %s disk publication failed (0x%x)\n",
                 ch->Index, drive ? "slave" : "master", st);
        dev->Present = false;
        AtaStopWorker(dev);
        IoDestroyRemoveLock(&dev->RemoveLock);
        return st;
    }

    DbgPrint("ATA: channel %u %s probe complete disk=%u\n",
             ch->Index, drive ? "slave" : "master", dev->DiskIndex);
    return NTLSTATUS_SUCCESS;
}

static KSTATUS AtaRegisterController(ATA_CONTROLLER* ctl) {
    if (!ctl) return NTLSTATUS_INVALID_PARAMETER;
    for (uint32_t i = 0; i < gAtaControllerCount; i++) {
        if (gAtaControllers[i] == ctl)
            return NTLSTATUS_SUCCESS;
    }
    if (gAtaControllerCount >= ATA_MAX_CONTROLLERS)
        return NTLSTATUS_NO_MEMORY;
    gAtaControllers[gAtaControllerCount++] = ctl;
    return NTLSTATUS_SUCCESS;
}

// Starts controller.
static KSTATUS AtaStartController(ATA_CONTROLLER* ctl) {
    if (!ctl || ctl->Magic != ATA_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("ATA: PciGetDeviceInfo failed: 0x%x\n", st);
        return st;
    }

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

    st = PciEnableDevice(ctl->Pdo, PCI_ENABLE_IO | PCI_ENABLE_INTX);
    if (!KSUCCESS(st)) {
        DbgPrint("ATA: PciEnableDevice failed: 0x%x\n", st);
        return st;
    }

    st = AtaRegisterController(ctl);
    if (!KSUCCESS(st))
        return st;

    for (uint8_t chIndex = 0; chIndex < ATA_CHANNEL_COUNT; chIndex++) {
        ATA_CHANNEL* ch = &ctl->Channels[chIndex];
        st = AtaInitChannel(ctl, chIndex);
        if (!KSUCCESS(st)) {
            DbgPrint("ATA: channel %u init failed (0x%x)\n", chIndex, st);
            continue;
        }

        st = AtaResetChannel(ch);
        if (!KSUCCESS(st))
            DbgPrint("ATA: channel %u reset did not complete cleanly (0x%x); probing drives with bounded waits\n",
                     chIndex, st);

        for (uint8_t drive = 0; drive < ATA_DRIVES_PER_CHANNEL; drive++) {
            DbgPrint("ATA: channel %u probing %s\n", chIndex, drive ? "slave" : "master");
            st = AtaProbeDrive(ctl, ch, drive);
            if (KSUCCESS(st)) {
                ctl->DeviceCount++;
            } else if (st != NTLSTATUS_NOT_FOUND) {
                DbgPrint("ATA: channel %u %s probe failed (0x%x)\n",
                         chIndex, drive ? "slave" : "master", st);
            }
        }
    }

    if (ctl->DeviceCount == 0)
        DbgPrint("ATA: no ATA disks found on PCI %02x:%02x.%x\n",
                 ctl->PciInfo.Address.Bus,
                 ctl->PciInfo.Address.Device,
                 ctl->PciInfo.Address.Function);

    ctl->Started = true;
    ctl->Fdo->PnpState = DevicePnpStarted;
    return NTLSTATUS_SUCCESS;
}

// Cancels release Remove Lock.
static void AtaCancelReleaseRemoveLock(IRP* irp, void* context) {
    ATA_DEVICE* dev = (ATA_DEVICE*)context;
    if (dev)
        IoReleaseRemoveLock(&dev->RemoveLock, irp);
}

// Releases controller Devices.
static void AtaRemoveControllerDevices(ATA_CONTROLLER* ctl, IO_DEVICE_CHANGE change) {
    if (!ctl) return;
    for (uint8_t c = 0; c < ATA_CHANNEL_COUNT; c++) {
        ATA_CHANNEL* ch = &ctl->Channels[c];
        for (uint8_t d = 0; d < ATA_DRIVES_PER_CHANNEL; d++) {
            ATA_DEVICE* dev = &ch->Devices[d];
            if (!dev->Present)
                continue;

            dev->Present = false;
            IrpQueueCancelAllWithCallback(&dev->IrpQueue, NTLSTATUS_OBJECT_DEAD, AtaCancelReleaseRemoveLock, dev);
            AtaStopWorker(dev);
            if (dev->DiskIndex != 0xFFu) {
                if (change == IoDeviceSurpriseRemoved)
                    IoNotifyDiskSurpriseRemoved(dev->DiskIndex);
                else
                    IoUnregisterDiskDevice(dev->DiskIndex);
                dev->DiskIndex = 0xFFu;
            }
        }
    }
}

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

    ATA_CONTROLLER* ctl = (ATA_CONTROLLER*)device->DeviceExtension;
    if (ctl->Magic == ATA_CONTROLLER_MAGIC) {
        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:
                return AtaStartController(ctl);

            case IRP_MN_QUERY_REMOVE_DEVICE:
            case IRP_MN_CANCEL_REMOVE_DEVICE:
                return NTLSTATUS_SUCCESS;

            case IRP_MN_SURPRISE_REMOVAL:
                ctl->SurpriseRemoved = true;
                device->PnpState = DevicePnpSurpriseRemoved;
                AtaRemoveControllerDevices(ctl, IoDeviceSurpriseRemoved);
                return NTLSTATUS_SUCCESS;

            case IRP_MN_REMOVE_DEVICE:
                ctl->Removed = true;
                device->PnpState = DevicePnpDeleted;
                AtaRemoveControllerDevices(ctl, IoDeviceRemoved);
                IoDetachDevice(device);
                IoDeleteDevice(device);
                return NTLSTATUS_SUCCESS;

            default:
                return IoForwardIrpToLower(device, irp);
        }
    }

    ATA_DEVICE* dev = (ATA_DEVICE*)device->DeviceExtension;
    if (dev->Magic == ATA_DEVICE_MAGIC) {
        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:
                return NTLSTATUS_SUCCESS;
            case IRP_MN_SURPRISE_REMOVAL:
                device->PnpState = DevicePnpSurpriseRemoved;
                dev->Present = false;
                IrpQueueCancelAllWithCallback(&dev->IrpQueue, NTLSTATUS_OBJECT_DEAD, AtaCancelReleaseRemoveLock, dev);
                AtaStopWorker(dev);
                return NTLSTATUS_SUCCESS;
            case IRP_MN_REMOVE_DEVICE:
                device->PnpState = DevicePnpDeleted;
                dev->Present = false;
                IrpQueueCancelAllWithCallback(&dev->IrpQueue, NTLSTATUS_OBJECT_DEAD, AtaCancelReleaseRemoveLock, dev);
                AtaStopWorker(dev);
                IoReleaseRemoveLockAndWait(&dev->RemoveLock, irp);
                IoDestroyRemoveLock(&dev->RemoveLock);
                IoDeleteDevice(device);
                dev->DeviceObject = (DEVICE_OBJECT*)0;
                return NTLSTATUS_SUCCESS;
            default:
                return NTLSTATUS_NOT_SUPPORTED;
        }
    }

    return NTLSTATUS_INVALID_PARAMETER;
}

KSTATUS AtaAddDevice(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(ATA_CONTROLLER), &fdo);
    if (!KSUCCESS(st))
        return st;
    (void)IoSetDeviceName(fdo, SYS_SV_LITERAL("ata-controller"));

    ATA_CONTROLLER* ctl = (ATA_CONTROLLER*)fdo->DeviceExtension;
    memset(ctl, 0, sizeof(*ctl));
    ctl->Magic = ATA_CONTROLLER_MAGIC;
    ctl->Fdo = fdo;
    ctl->Pdo = physicalDeviceObject;

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

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

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

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

    gAtaControllerCount = 0;

    memset(&gAtaDriver, 0, sizeof(gAtaDriver));
    IoInitializeDriverObject(&gAtaDriver, SYS_SV_LITERAL("ata"));
    gAtaDriver.IdTable = &gAtaIdTable;
    gAtaDriver.DriverAddDevice = AtaAddDevice;
    gAtaDriver.MajorFunction[IRP_MJ_READ] = AtaDispatchReadWrite;
    gAtaDriver.MajorFunction[IRP_MJ_WRITE] = AtaDispatchReadWrite;
    gAtaDriver.MajorFunction[IRP_MJ_FLUSH_BUFFERS] = AtaDispatchFlush;
    gAtaDriver.MajorFunction[IRP_MJ_PNP] = AtaDispatchPnp;
    return NTLSTATUS_SUCCESS;
}

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