// ATA compatibility storage driver. This file owns legacy IDE probing and PIO disk access for older virtual hardware.

#include "ata.h"

const NTL_DDK* gAtaDdk;
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]),
};

// Implements the in8 operation.
static uint8_t AtaIn8(ATA_CHANNEL* ch, uint8_t reg) {
    return __inbyte((unsigned short)(ch->CommandBase + reg));
}

// Implements the in16 operation.
static uint16_t AtaIn16(ATA_CHANNEL* ch, uint8_t reg) {
    return __inword((unsigned short)(ch->CommandBase + reg));
}

// Implements the out8 operation.
static void AtaOut8(ATA_CHANNEL* ch, uint8_t reg, uint8_t value) {
    __outbyte((unsigned short)(ch->CommandBase + reg), value);
}

// Implements the out16 operation.
static void AtaOut16(ATA_CHANNEL* ch, uint8_t reg, uint16_t value) {
    __outword((unsigned short)(ch->CommandBase + reg), value);
}

// Implements the alt Status operation.
static uint8_t AtaAltStatus(ATA_CHANNEL* ch) {
    return __inbyte((unsigned short)(ch->ControlBase + ATA_REG_ALT_STATUS));
}

// Implements the device Control operation.
static void AtaDeviceControl(ATA_CHANNEL* ch, uint8_t value) {
    __outbyte((unsigned short)(ch->ControlBase + ATA_REG_DEVICE_CONTROL), value);
}

// Implements the delay400ns operation.
static void AtaDelay400ns(ATA_CHANNEL* ch) {
    (void)AtaAltStatus(ch);
    (void)AtaAltStatus(ch);
    (void)AtaAltStatus(ch);
    (void)AtaAltStatus(ch);
}

// Waits for not Busy.
static KSTATUS AtaWaitNotBusy(ATA_CHANNEL* ch) {
    for (uint32_t i = 0; i < ATA_POLL_LIMIT; i++) {
        if (!(AtaIn8(ch, ATA_REG_STATUS) & ATA_STATUS_BSY))
            return NTLSTATUS_SUCCESS;
        KeYield();
    }
    return NTLSTATUS_TIMEOUT;
}

// Waits for drq.
static KSTATUS AtaWaitDrq(ATA_CHANNEL* ch) {
    for (uint32_t i = 0; i < ATA_POLL_LIMIT; i++) {
        uint8_t st = AtaIn8(ch, ATA_REG_STATUS);
        if (st & ATA_STATUS_ERR)
            return NTLSTATUS_DEVICE_ERROR;
        if (st & ATA_STATUS_DF)
            return NTLSTATUS_DEVICE_ERROR;
        if (!(st & ATA_STATUS_BSY) && (st & ATA_STATUS_DRQ))
            return NTLSTATUS_SUCCESS;
        KeYield();
    }
    return NTLSTATUS_TIMEOUT;
}

// Implements the select Drive operation.
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 the identify operation.
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);
    if (!KSUCCESS(st))
        return st;

    AtaSelectDrive(dev);
    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 = AtaIn8(ch, ATA_REG_STATUS);
    if (status == 0) {
        KeSignalSemaphore(ch->Lock);
        return NTLSTATUS_NOT_FOUND;
    }

    st = AtaWaitNotBusy(ch);
    if (!KSUCCESS(st)) {
        KeSignalSemaphore(ch->Lock);
        return st;
    }

    if (AtaIn8(ch, ATA_REG_LBA1) != 0 || AtaIn8(ch, ATA_REG_LBA2) != 0) {
        KeSignalSemaphore(ch->Lock);
        return NTLSTATUS_NOT_SUPPORTED; // ATAPI or non-ATA signature.
    }

    st = AtaWaitDrq(ch);
    if (!KSUCCESS(st)) {
        KeSignalSemaphore(ch->Lock);
        return st;
    }

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

    KeSignalSemaphore(ch->Lock);

    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];
    }

    return dev->SectorCount ? NTLSTATUS_SUCCESS : NTLSTATUS_NOT_FOUND;
}

// Implements the program Lba operation.
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);
    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);
}

// 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);
    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;
}

// Implements the worker operation.
static void AtaWorker(uintptr_t arg) {
    ATA_DEVICE* dev = (ATA_DEVICE*)arg;
    if (!dev)
        return;

    while (true) {
        KeWaitSemaphore(dev->IrpQueue.Semaphore);

        IRP* irp = IrpQueueDequeue(&dev->IrpQueue);
        if (!irp)
            continue;

        bool write = false;
        uint64_t lba = 0;
        uint32_t sectors = 0;
        void* buffer = (void*)0;
        KSTATUS 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);
        }

        IoReleaseRemoveLock(&dev->RemoveLock, irp);
        IoCompleteRequest(irp, st, KSUCCESS(st) ? ((size_t)sectors * ATA_SECTOR_SIZE) : 0);
    }
}

// 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;
}

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

// Implements the bar Io Base operation.
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, 0x02u); // 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);

    dev->Present = true;
    st = AtaRegisterDisk(dev);
    if (!KSUCCESS(st)) {
        dev->Present = false;
        return st;
    }

    st = KeCreateKernelThread((uintptr_t)AtaWorker, (uintptr_t)dev,
                              PriorityHigh, &dev->WorkerThread);
    if (!KSUCCESS(st))
        DbgPrint("ATA: drive %u worker thread create failed (0x%x)\n", dev->DiskIndex, st);
    return NTLSTATUS_SUCCESS;
}

// Registers controller.
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;
        }

        for (uint8_t drive = 0; drive < ATA_DRIVES_PER_CHANNEL; drive++) {
            st = AtaProbeDrive(ctl, ch, drive);
            if (KSUCCESS(st))
                ctl->DeviceCount++;
        }
    }

    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;

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

// 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;
                IrpQueueCancelAllEx(&dev->IrpQueue, NTLSTATUS_OBJECT_DEAD, AtaCancelReleaseRemoveLock, dev);
                return NTLSTATUS_SUCCESS;
            case IRP_MN_REMOVE_DEVICE:
                device->PnpState = DevicePnpDeleted;
                dev->Present = false;
                IrpQueueCancelAllEx(&dev->IrpQueue, NTLSTATUS_OBJECT_DEAD, AtaCancelReleaseRemoveLock, dev);
                IoReleaseRemoveLockAndWait(&dev->RemoveLock, irp);
                IoDeleteDevice(device);
                dev->DeviceObject = (DEVICE_OBJECT*)0;
                return NTLSTATUS_SUCCESS;
            default:
                return NTLSTATUS_NOT_SUPPORTED;
        }
    }

    return NTLSTATUS_INVALID_PARAMETER;
}

// Implements the add Device operation.
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;
}

// Implements the driver Entry operation.
KSTATUS DriverEntry(const NTL_DDK* ddk, NTL_DRIVER_INIT* init) {
    (void)init;
    if (!ddk)
        return NTLSTATUS_INVALID_PARAMETER;

    DbgPrint("ATA: driver loading...\n");
    gAtaDdk = ddk;
    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_PNP] = AtaDispatchPnp;

    KSTATUS st = IoRegisterDriverObject(&gAtaDriver);
    if (!KSUCCESS(st))
        DbgPrint("ATA: IoRegisterDriverObject failed: 0x%x\n", st);
    return st;
}