// I/O manager. This file owns driver objects, device stacks, PnP dispatch, and the disk namespace.

#include <string.h>
#include <intrin.h>
#include <ntlite/ddi/io_driver.h>
#include <kernel/ke_irp.h>
#include <kernel/mm_virt.h>
#include <kernel/ke_debug.h>
#include <kernel/ob_object.h>
#include <kernel/ke_sched.h>

#define MAX_DRIVER_OBJECTS    64
#define MAX_PDO_OBJECTS       128

// Implements the sv Or Unnamed operation.
static SYS_STRING_VIEW IoSvOrUnnamed(SYS_STRING_VIEW name) {
    return SysSvIsEmpty(name) ? SYS_SV_LITERAL("<unnamed>") : name;
}

// Implements the bus Type Name operation.
static const char* IoBusTypeName(IO_BUS_TYPE busType) {
    switch (busType) {
        case IoBusPci: return "pci";
        case IoBusUsb: return "usb";
        default: return "unknown";
    }
}

// Driver object registry / PDO AddDevice binding.

static DRIVER_OBJECT* gDriverObjects[MAX_DRIVER_OBJECTS];
static uint32_t gDriverObjectCount;
static DEVICE_OBJECT* gPhysicalDevices[MAX_PDO_OBJECTS];
static uint32_t gPhysicalDeviceCount;

// Implements the driver Matches Pdo operation.
static bool IoDriverMatchesPdo(const DRIVER_OBJECT* driver, const DEVICE_OBJECT* pdo) {
    if (!driver || !driver->IdTable || !pdo || pdo->Type != DeviceObjectPhysical)
        return false;

    for (size_t d = 0; d < driver->IdTable->Count; d++) {
        const IO_DEVICE_ID* want = &driver->IdTable->Ids[d];
        if (want->BusType != IoBusUnknown && want->BusType != pdo->u.Pdo.BusType)
            continue;
        for (uint8_t p = 0; p < pdo->u.Pdo.IdCount; p++) {
            if (SysSvEqual(want->Id, pdo->u.Pdo.Ids[p].Id))
                return true;
        }
    }
    return false;
}

// Implements the try Add Device For Driver operation.
static bool IoTryAddDeviceForDriver(DRIVER_OBJECT* driver, DEVICE_OBJECT* pdo) {
    if (!driver || !pdo || !driver->DriverAddDevice)
        return false;
    if (!IoDriverMatchesPdo(driver, pdo))
        return false;

    KSTATUS st = driver->DriverAddDevice(driver, pdo);
    if (KSUCCESS(st)) {
        DbgPrint("io: AddDevice driver=%v pdo=%v bus=%s\n",
                 IoSvOrUnnamed(driver->Name), IoSvOrUnnamed(pdo->Name), IoBusTypeName(pdo->u.Pdo.BusType));

        KSTATUS start = IoSendPnpIrpSync(pdo, IRP_MN_START_DEVICE);
        if (!KSUCCESS(start) && start != STATUS_NOT_SUPPORTED) {
            DbgPrint("io: START_DEVICE failed driver=%v pdo=%v st=0x%x\n",
                     IoSvOrUnnamed(driver->Name), IoSvOrUnnamed(pdo->Name), start);
        }
        return true;
    } else {
        DbgPrint("io: AddDevice failed driver=%v pdo=%v st=0x%x\n",
                 IoSvOrUnnamed(driver->Name), IoSvOrUnnamed(pdo->Name), st);
        return true;
    }
}

__declspec(dllexport) void IoInitializeDriverObject(DRIVER_OBJECT* driver,
                                                    SYS_STRING_VIEW name) {
    if (!driver) return;
    memset(driver, 0, sizeof(*driver));
    driver->Header.Type = ObjectTypeDriver;
    driver->Header.RefCount = 1;
    driver->Header.Destroy = (void*)0;
    driver->Name = name;
}

// Registers driver Object.
__declspec(dllexport) KSTATUS IoRegisterDriverObject(DRIVER_OBJECT* driver) {
    if (!driver) return STATUS_INVALID_PARAMETER;
    if (gDriverObjectCount >= MAX_DRIVER_OBJECTS)
        return STATUS_NO_MEMORY;

    if (driver->Header.Type != ObjectTypeDriver) {
        driver->Header.Type = ObjectTypeDriver;
        driver->Header.RefCount = 1;
    }

    gDriverObjects[gDriverObjectCount++] = driver;
    DbgPrint("io: registered driver object %v\n", IoSvOrUnnamed(driver->Name));

    for (uint32_t i = 0; i < gPhysicalDeviceCount; i++)
        (void)IoTryAddDeviceForDriver(driver, gPhysicalDevices[i]);

    return STATUS_SUCCESS;
}

// Device object lifetime.

static void IoDeviceDestroy(OBJECT* object) {
    DEVICE_OBJECT* dev = (DEVICE_OBJECT*)object;

    for (uint32_t i = 0; i < gPhysicalDeviceCount; i++) {
        if (gPhysicalDevices[i] == dev) {
            gPhysicalDevices[i] = gPhysicalDevices[gPhysicalDeviceCount - 1u];
            gPhysicalDevices[gPhysicalDeviceCount - 1u] = (DEVICE_OBJECT*)0;
            gPhysicalDeviceCount--;
            break;
        }
    }

    if (dev->LowerDevice)
        IoDetachDevice(dev);

    if (dev->Parent) {
        DEVICE_OBJECT* parent = dev->Parent;
        uintptr_t flags = KeAcquireSpinlock(&parent->Lock);
        DEVICE_OBJECT** link = &parent->FirstChild;
        while (*link) {
            if (*link == dev) {
                *link = dev->NextSibling;
                break;
            }
            link = &(*link)->NextSibling;
        }
        KeReleaseSpinlock(&parent->Lock, flags);
        dev->Parent = (DEVICE_OBJECT*)0;
        dev->NextSibling = (DEVICE_OBJECT*)0;
        IoDereferenceDeviceObject(parent);
    }
}

static KSTATUS IoCreateDeviceInternal(DRIVER_OBJECT* driver,
                                      DEVICE_OBJECT_TYPE type,
                                      size_t extensionSize,
                                      DEVICE_OBJECT** outDevice) {
    if (!driver || !outDevice) return STATUS_INVALID_PARAMETER;

    size_t total = sizeof(DEVICE_OBJECT) + extensionSize;
    DEVICE_OBJECT* dev = (DEVICE_OBJECT*)ObAllocateObject(ObjectTypeDevice, total);
    if (!dev) return STATUS_NO_MEMORY;

    dev->Header.Destroy = IoDeviceDestroy;
    dev->DriverObject = driver;
    dev->Type = type;
    dev->PnpState = DevicePnpNotStarted;
    dev->AlignmentRequirement = 1;
    dev->Lock = (SPINLOCK)SPINLOCK_INIT;

    // The per-device union is interpreted differently for PDOs and disk FDOs.
    // Do not write disk defaults into a PDO: u.Disk.DiskIndex aliases the low
    // byte of u.Pdo.BusType, which prevents ID-table AddDevice matching.
    if (type != DeviceObjectPhysical) {
        dev->u.Disk.DiskIndex = 0xFF;
        dev->u.Disk.SectorSize = 512;
    }

    if (extensionSize)
        dev->DeviceExtension = (void*)((uint8_t*)dev + sizeof(DEVICE_OBJECT));

    *outDevice = dev;
    return STATUS_SUCCESS;
}

__declspec(dllexport) KSTATUS IoCreateDevice(DRIVER_OBJECT* driver,
                                             size_t extensionSize,
                                             DEVICE_OBJECT** outDevice) {
    return IoCreateDeviceInternal(driver, DeviceObjectFunction, extensionSize, outDevice);
}

__declspec(dllexport) KSTATUS IoCreatePhysicalDevice(DRIVER_OBJECT* busDriver,
                                                     DEVICE_OBJECT* parent,
                                                     const IO_DEVICE_ID* ids,
                                                     size_t idCount,
                                                     size_t extensionSize,
                                                     DEVICE_OBJECT** outPdo) {
    if (!busDriver || !outPdo) return STATUS_INVALID_PARAMETER;
    if (idCount > IO_MAX_DEVICE_IDS) return STATUS_INVALID_PARAMETER;

    DEVICE_OBJECT* pdo = (DEVICE_OBJECT*)0;
    KSTATUS st = IoCreateDeviceInternal(busDriver, DeviceObjectPhysical, extensionSize, &pdo);
    if (!KSUCCESS(st)) return st;

    if (idCount && ids) {
        for (size_t i = 0; i < idCount; i++) {
            if (!SysSvIsValid(ids[i].Id) || ids[i].Id.Length >= IO_MAX_DEVICE_ID_CHARS) {
                IoDeleteDevice(pdo);
                return STATUS_INVALID_PARAMETER;
            }

            char* dst = pdo->u.Pdo.IdStorage[i];
            if (ids[i].Id.Length)
                memcpy(dst, ids[i].Id.Data, ids[i].Id.Length);
            dst[ids[i].Id.Length] = '\0';

            pdo->u.Pdo.Ids[i].BusType = ids[i].BusType;
            pdo->u.Pdo.Ids[i].Id = SysSvFromParts(dst, ids[i].Id.Length);
            if (pdo->u.Pdo.BusType == IoBusUnknown)
                pdo->u.Pdo.BusType = ids[i].BusType;
        }
        pdo->u.Pdo.IdCount = (uint8_t)idCount;
    }

    if (parent) {
        IoReferenceDeviceObject(parent);
        pdo->Parent = parent;
        uintptr_t flags = KeAcquireSpinlock(&parent->Lock);
        pdo->NextSibling = parent->FirstChild;
        parent->FirstChild = pdo;
        KeReleaseSpinlock(&parent->Lock, flags);
    }

    *outPdo = pdo;
    return STATUS_SUCCESS;
}

// Registers physical Device.
__declspec(dllexport) KSTATUS IoRegisterPhysicalDevice(DEVICE_OBJECT* pdo) {
    if (!pdo || pdo->Type != DeviceObjectPhysical)
        return STATUS_INVALID_PARAMETER;

    for (uint32_t i = 0; i < gPhysicalDeviceCount; i++) {
        if (gPhysicalDevices[i] == pdo)
            return STATUS_ALREADY_EXISTS;
    }

    if (gPhysicalDeviceCount >= MAX_PDO_OBJECTS)
        return STATUS_NO_MEMORY;

    gPhysicalDevices[gPhysicalDeviceCount++] = pdo;

    bool anyMatched = false;
    for (uint32_t i = 0; i < gDriverObjectCount; i++) {
        if (IoTryAddDeviceForDriver(gDriverObjects[i], pdo))
            anyMatched = true;
    }

    if (!anyMatched && pdo->u.Pdo.BusType == IoBusUsb) {
        SYS_STRING_VIEW id = pdo->u.Pdo.IdCount ? pdo->u.Pdo.Ids[0].Id : SYS_SV_LITERAL("<no-id>");
        DbgPrint("io: no function driver for pdo=%v id=%v bus=%s\n",
                 IoSvOrUnnamed(pdo->Name), id, IoBusTypeName(pdo->u.Pdo.BusType));
    }
    return STATUS_SUCCESS;
}


__declspec(dllexport) KSTATUS IoSetDeviceName(DEVICE_OBJECT* device,
                                               SYS_STRING_VIEW name) {
    if (!device || !SysSvIsValid(name))
        return STATUS_INVALID_PARAMETER;
    if (name.Length >= IO_MAX_DEVICE_NAME_CHARS)
        return STATUS_BUFFER_TOO_SMALL;

    if (name.Length)
        memcpy(device->NameStorage, name.Data, name.Length);
    device->NameStorage[name.Length] = '\0';
    device->Name = SysSvFromParts(device->NameStorage, name.Length);
    return STATUS_SUCCESS;
}

// Implements the reference Device Object operation.
__declspec(dllexport) void IoReferenceDeviceObject(DEVICE_OBJECT* device) {
    if (device) ObReferenceObject(&device->Header);
}

// Implements the dereference Device Object operation.
__declspec(dllexport) void IoDereferenceDeviceObject(DEVICE_OBJECT* device) {
    if (device) ObDereferenceObject(&device->Header);
}

// Releases physical Device Registry Entry.
static void IoRemovePhysicalDeviceRegistryEntry(DEVICE_OBJECT* device) {
    if (!device || device->Type != DeviceObjectPhysical)
        return;

    for (uint32_t i = 0; i < gPhysicalDeviceCount; i++) {
        if (gPhysicalDevices[i] == device) {
            for (uint32_t j = i + 1; j < gPhysicalDeviceCount; j++)
                gPhysicalDevices[j - 1] = gPhysicalDevices[j];
            gPhysicalDeviceCount--;
            if (gPhysicalDeviceCount < MAX_PDO_OBJECTS)
                gPhysicalDevices[gPhysicalDeviceCount] = (DEVICE_OBJECT*)0;
            return;
        }
    }
}

// Implements the unlink Device From Parent operation.
static void IoUnlinkDeviceFromParent(DEVICE_OBJECT* device) {
    if (!device || !device->Parent)
        return;

    DEVICE_OBJECT* parent = device->Parent;
    uintptr_t flags = KeAcquireSpinlock(&parent->Lock);
    DEVICE_OBJECT** link = &parent->FirstChild;
    while (*link) {
        if (*link == device) {
            *link = device->NextSibling;
            break;
        }
        link = &(*link)->NextSibling;
    }
    device->Parent = (DEVICE_OBJECT*)0;
    device->NextSibling = (DEVICE_OBJECT*)0;
    KeReleaseSpinlock(&parent->Lock, flags);

    IoDereferenceDeviceObject(parent);
}

// Releases device.
__declspec(dllexport) void IoDeleteDevice(DEVICE_OBJECT* device) {
    if (!device) return;
    if (device->DeletePending)
        return;

    device->DeletePending = true;
    device->PnpState = DevicePnpDeleted;
    IoUnlinkDeviceFromParent(device);
    IoRemovePhysicalDeviceRegistryEntry(device);
    IoDereferenceDeviceObject(device);
}

// Device stack management.

__declspec(dllexport) DEVICE_OBJECT* IoAttachDeviceToDeviceStack(DEVICE_OBJECT* source,
                                                                 DEVICE_OBJECT* target) {
    if (!source || !target) return (DEVICE_OBJECT*)0;

    DEVICE_OBJECT* top = target;
    while (top->AttachedDevice)
        top = top->AttachedDevice;

    top->AttachedDevice = source;
    source->LowerDevice = top;
    IoReferenceDeviceObject(top);
    return top;
}

// Implements the detach Device operation.
__declspec(dllexport) void IoDetachDevice(DEVICE_OBJECT* source) {
    if (!source || !source->LowerDevice) return;

    DEVICE_OBJECT* lower = source->LowerDevice;
    if (lower->AttachedDevice == source)
        lower->AttachedDevice = (DEVICE_OBJECT*)0;
    source->LowerDevice = (DEVICE_OBJECT*)0;
    IoDereferenceDeviceObject(lower);
}

// Returns attached Device Reference.
__declspec(dllexport) DEVICE_OBJECT* IoGetAttachedDeviceReference(DEVICE_OBJECT* device) {
    if (!device) return (DEVICE_OBJECT*)0;
    DEVICE_OBJECT* top = device;
    while (top->AttachedDevice)
        top = top->AttachedDevice;
    IoReferenceDeviceObject(top);
    return top;
}

// Device-stack helpers / PnP dispatch.

__declspec(dllexport) uint8_t IoGetDeviceStackSize(DEVICE_OBJECT* device) {
    uint8_t count = 0;
    for (DEVICE_OBJECT* cur = device; cur && count < 0xFEu; cur = cur->LowerDevice)
        count++;
    return count ? count : 1;
}

// Implements the forward Irp To Lower operation.
__declspec(dllexport) KSTATUS IoForwardIrpToLower(DEVICE_OBJECT* device, IRP* irp) {
    if (!device || !irp || !device->LowerDevice)
        return STATUS_NOT_SUPPORTED;

    IO_STACK_LOCATION* cur = IoGetCurrentIrpStackLocation(irp);
    IO_STACK_LOCATION* next = IoGetNextIrpStackLocation(irp);
    if (!cur || !next)
        return STATUS_INVALID_PARAMETER;

    DEVICE_OBJECT* currentDevice = cur->DeviceObject;
    memcpy(next, cur, sizeof(*next));

    if (currentDevice == device) {
        cur->DeviceObject = (DEVICE_OBJECT*)0;
        _InterlockedDecrement(&device->OutstandingIrps);
        IoDereferenceDeviceObject(device);
    }

    irp->CurrentLocation++;
    return IoCallDriver(device->LowerDevice, irp);
}

__declspec(dllexport) KSTATUS IoSendPnpIrpSync(DEVICE_OBJECT* device,
                                                IO_PNP_MINOR_FUNCTION minor) {
    if (!device) return STATUS_INVALID_PARAMETER;

    DEVICE_OBJECT* top = IoGetAttachedDeviceReference(device);
    if (!top) return STATUS_INVALID_PARAMETER;

    uint8_t stackCount = IoGetDeviceStackSize(top);
    IRP* irp = (IRP*)0;
    KSTATUS st = IoAllocateIrp(stackCount, &irp);
    if (!KSUCCESS(st)) {
        IoDereferenceDeviceObject(top);
        return st;
    }

    IoBuildPnpIrp(irp, minor);
    st = IoSubmitIrpSync(top, irp);
    IoFreeIrp(irp);
    IoDereferenceDeviceObject(top);
    return st;
}

// Disk namespace.

#define MAX_DISK_DEVICES 8
#define MAX_DISK_CALLBACKS 16

typedef struct IO_DISK_CHANGE_CALLBACK_ENTRY {
    bool Used;
    IO_DISK_CHANGE_CALLBACK Callback;
    void* Context;
} IO_DISK_CHANGE_CALLBACK_ENTRY;

typedef struct IO_DISK_ARRIVAL_CALLBACK_ENTRY {
    bool Used;
    IO_DISK_ARRIVAL_CALLBACK Callback;
    void* Context;
} IO_DISK_ARRIVAL_CALLBACK_ENTRY;

static DEVICE_OBJECT* gDiskDevices[MAX_DISK_DEVICES];
static IO_DISK_CHANGE_CALLBACK_ENTRY gDiskChangeCallbacks[MAX_DISK_CALLBACKS];
static uint32_t gDiskChangeCallbackCount;
static IO_DISK_ARRIVAL_CALLBACK_ENTRY gDiskArrivalCallbacks[MAX_DISK_CALLBACKS];
static uint32_t gDiskArrivalCallbackCount;

#define IO_DISK_READAHEAD_MAX_SECTORS 128u
#define IO_DISK_READAHEAD_TRIGGER_SCORE 2u
#define IO_DISK_READAHEAD_MIN_SECTORS 8u
#define IO_DISK_SECTOR_BYTES 512u

typedef struct IO_DISK_READAHEAD_STATE {
    SPINLOCK Lock;
    uint8_t* Buffer;
    uint8_t  Valid;
    uint8_t  Loading;
    uint64_t Generation;
    uint64_t BaseLba;
    uint32_t Count;
    uint64_t NextLba;
    uint32_t SequentialScore;
    IO_DISK_STATS Stats;
} IO_DISK_READAHEAD_STATE;

static IO_DISK_READAHEAD_STATE gDiskReadAhead[MAX_DISK_DEVICES];

// Updates a maximum statistic while holding the disk read-ahead lock.
static void IoDiskStatMaxLocked(uint64_t* field, uint64_t value) {
    if (field && *field < value) *field = value;
}

// Prepares the per-disk read-ahead buffer.
static void IoDiskReadAheadPrepare(uint8_t driveIndex) {
    if (driveIndex >= MAX_DISK_DEVICES)
        return;

    IO_DISK_READAHEAD_STATE* state = &gDiskReadAhead[driveIndex];
    state->Lock = (SPINLOCK)SPINLOCK_INIT;
    state->Valid = 0;
    state->Loading = 0;
    state->Generation = 0;
    state->BaseLba = 0;
    state->Count = 0;
    state->NextLba = 0;
    state->SequentialScore = 0;
    memset(&state->Stats, 0, sizeof(state->Stats));

    if (!state->Buffer) {
        state->Buffer = (uint8_t*)MmAllocatePoolTag(IO_DISK_READAHEAD_MAX_SECTORS * IO_DISK_SECTOR_BYTES,
                                                    MM_POOL_TAG('d', 'r', 'a', 'h'));
        if (!state->Buffer)
            state->Stats.ReadaheadAllocFailures++;
    }
}

// Releases the per-disk read-ahead buffer.
static void IoDiskReadAheadRelease(uint8_t driveIndex) {
    if (driveIndex >= MAX_DISK_DEVICES)
        return;

    IO_DISK_READAHEAD_STATE* state = &gDiskReadAhead[driveIndex];
    uint8_t* buffer = (uint8_t*)0;
    for (;;) {
        uintptr_t flags = KeAcquireSpinlock(&state->Lock);
        if (!state->Loading) {
            buffer = state->Buffer;
            state->Buffer = (uint8_t*)0;
            state->Valid = 0;
            state->Generation++;
            state->Count = 0;
            state->SequentialScore = 0;
            KeReleaseSpinlock(&state->Lock, flags);
            break;
        }
        KeReleaseSpinlock(&state->Lock, flags);
        KeYield();
    }

    if (buffer)
        MmFreePool(buffer);
}

// Invalidates cached disk read-ahead data for one disk.
static void IoDiskReadAheadInvalidate(uint8_t driveIndex) {
    if (driveIndex >= MAX_DISK_DEVICES)
        return;

    IO_DISK_READAHEAD_STATE* state = &gDiskReadAhead[driveIndex];
    uintptr_t flags = KeAcquireSpinlock(&state->Lock);
    state->Valid = 0;
    state->Generation++;
    state->Count = 0;
    state->Stats.ReadaheadInvalidations++;
    KeReleaseSpinlock(&state->Lock, flags);
}

// Returns true when one range fully contains another range.
static bool IoDiskRangeContains(uint64_t outerBase, uint32_t outerCount,
                                uint64_t innerBase, uint32_t innerCount) {
    uint64_t outerEnd = outerBase + outerCount;
    uint64_t innerEnd = innerBase + innerCount;
    if (outerEnd < outerBase || innerEnd < innerBase)
        return false;
    return innerBase >= outerBase && innerEnd <= outerEnd;
}

// Attempts to satisfy a disk read from the per-disk read-ahead buffer.
static bool IoDiskReadAheadTryHit(uint8_t driveIndex, uint64_t lba, uint32_t count, void* buffer) {
    if (driveIndex >= MAX_DISK_DEVICES || !buffer || count == 0)
        return false;

    IO_DISK_READAHEAD_STATE* state = &gDiskReadAhead[driveIndex];
    uintptr_t flags = KeAcquireSpinlock(&state->Lock);
    bool hit = state->Buffer && state->Valid && !state->Loading &&
               IoDiskRangeContains(state->BaseLba, state->Count, lba, count);
    if (hit) {
        size_t offset = (size_t)(lba - state->BaseLba) * IO_DISK_SECTOR_BYTES;
        size_t bytes = (size_t)count * IO_DISK_SECTOR_BYTES;
        memcpy(buffer, state->Buffer + offset, bytes);
        state->Stats.ReadaheadHits++;
        state->Stats.ReadaheadBytesServed += bytes;
    } else {
        state->Stats.ReadaheadMisses++;
    }
    KeReleaseSpinlock(&state->Lock, flags);
    return hit;
}

// Records request size and sequential access state for one disk read.
static uint32_t IoDiskReadAheadNoteRead(uint8_t driveIndex, uint64_t lba, uint32_t count) {
    if (driveIndex >= MAX_DISK_DEVICES || count == 0)
        return 0;

    IO_DISK_READAHEAD_STATE* state = &gDiskReadAhead[driveIndex];
    uintptr_t flags = KeAcquireSpinlock(&state->Lock);
    state->Stats.ReadRequests++;
    state->Stats.ReadSectors += count;
    IoDiskStatMaxLocked(&state->Stats.MaxReadSectors, count);

    if (lba == state->NextLba) {
        if (state->SequentialScore < 32)
            state->SequentialScore++;
    } else {
        state->SequentialScore = 1;
    }
    state->NextLba = lba + count;
    uint32_t score = state->SequentialScore;
    if (score >= IO_DISK_READAHEAD_TRIGGER_SCORE)
        state->Stats.ReadaheadSequentialDetections++;
    KeReleaseSpinlock(&state->Lock, flags);
    return score;
}

// Records one disk write and invalidates stale read-ahead data.
static void IoDiskReadAheadNoteWrite(uint8_t driveIndex, uint32_t count) {
    if (driveIndex >= MAX_DISK_DEVICES)
        return;

    IO_DISK_READAHEAD_STATE* state = &gDiskReadAhead[driveIndex];
    uintptr_t flags = KeAcquireSpinlock(&state->Lock);
    state->Stats.WriteRequests++;
    state->Stats.WriteSectors += count;
    IoDiskStatMaxLocked(&state->Stats.MaxWriteSectors, count);
    state->Valid = 0;
    state->Generation++;
    state->Count = 0;
    state->Stats.ReadaheadInvalidations++;
    KeReleaseSpinlock(&state->Lock, flags);
}

// Implements the notify Disk Change operation.
static void IoNotifyDiskChange(uint8_t driveIndex, DEVICE_OBJECT* device, IO_DEVICE_CHANGE change) {
    for (uint32_t i = 0; i < gDiskChangeCallbackCount; i++) {
        IO_DISK_CHANGE_CALLBACK_ENTRY* entry = &gDiskChangeCallbacks[i];
        if (entry->Used && entry->Callback)
            entry->Callback(driveIndex, device, change, entry->Context);
    }
    if (change == IoDeviceArrived) {
        for (uint32_t i = 0; i < gDiskArrivalCallbackCount; i++) {
            IO_DISK_ARRIVAL_CALLBACK_ENTRY* entry = &gDiskArrivalCallbacks[i];
            if (entry->Used && entry->Callback)
                entry->Callback(driveIndex, device, entry->Context);
        }
    }
}

__declspec(dllexport) KSTATUS IoRegisterDiskDevice(uint8_t driveIndex,
                                                   DEVICE_OBJECT* device) {
    if (driveIndex >= MAX_DISK_DEVICES || !device || !device->DriverObject)
        return STATUS_INVALID_PARAMETER;
    if (gDiskDevices[driveIndex])
        return STATUS_ALREADY_EXISTS;

    device->Flags |= DEVICE_OBJECT_DISK;
    device->u.Disk.DiskIndex = driveIndex;
    if (device->u.Disk.SectorSize == 0)
        device->u.Disk.SectorSize = 512;

    KSTATUS pnpStatus = IoSendPnpIrpSync(device, IRP_MN_START_DEVICE);
    if (!KSUCCESS(pnpStatus) && pnpStatus != STATUS_NOT_SUPPORTED)
        return pnpStatus;
    if (pnpStatus == STATUS_NOT_SUPPORTED)
        device->PnpState = DevicePnpStarted;

    IoDiskReadAheadPrepare(driveIndex);

    IoReferenceDeviceObject(device);
    gDiskDevices[driveIndex] = device;
    IoNotifyDiskChange(driveIndex, device, IoDeviceArrived);
    return STATUS_SUCCESS;
}

// Implements the notify Disk Surprise Removed operation.
__declspec(dllexport) KSTATUS IoNotifyDiskSurpriseRemoved(uint8_t driveIndex) {
    if (driveIndex >= MAX_DISK_DEVICES)
        return STATUS_INVALID_PARAMETER;
    DEVICE_OBJECT* dev = gDiskDevices[driveIndex];
    if (!dev) return STATUS_NOT_FOUND;

    IoDiskReadAheadInvalidate(driveIndex);

    dev->PnpState = DevicePnpSurpriseRemoved;
    IoNotifyDiskChange(driveIndex, dev, IoDeviceSurpriseRemoved);
    (void)IoSendPnpIrpSync(dev, IRP_MN_SURPRISE_REMOVAL);
    return STATUS_SUCCESS;
}

// Implements the unregister Disk Device operation.
__declspec(dllexport) KSTATUS IoUnregisterDiskDevice(uint8_t driveIndex) {
    if (driveIndex >= MAX_DISK_DEVICES)
        return STATUS_INVALID_PARAMETER;
    DEVICE_OBJECT* dev = gDiskDevices[driveIndex];
    if (!dev) return STATUS_NOT_FOUND;

    if (dev->PnpState != DevicePnpSurpriseRemoved) {
        IoNotifyDiskChange(driveIndex, dev, IoDeviceQueryRemove);
        KSTATUS queryStatus = IoSendPnpIrpSync(dev, IRP_MN_QUERY_REMOVE_DEVICE);
        if (!KSUCCESS(queryStatus) && queryStatus != STATUS_NOT_SUPPORTED)
            return queryStatus;

        dev->PnpState = DevicePnpSurpriseRemoved;
        IoNotifyDiskChange(driveIndex, dev, IoDeviceSurpriseRemoved);
        (void)IoSendPnpIrpSync(dev, IRP_MN_SURPRISE_REMOVAL);
    }

    dev->PnpState = DevicePnpRemovePending;
    (void)IoSendPnpIrpSync(dev, IRP_MN_REMOVE_DEVICE);

    gDiskDevices[driveIndex] = (DEVICE_OBJECT*)0;
    IoDiskReadAheadRelease(driveIndex);
    IoNotifyDiskChange(driveIndex, dev, IoDeviceRemoved);
    IoDereferenceDeviceObject(dev);
    return STATUS_SUCCESS;
}

__declspec(dllexport) KSTATUS IoRegisterDiskArrivalCallback(IO_DISK_ARRIVAL_CALLBACK callback,
                                                            void* context) {
    if (!callback)
        return STATUS_INVALID_PARAMETER;
    if (gDiskArrivalCallbackCount >= MAX_DISK_CALLBACKS)
        return STATUS_NO_MEMORY;

    IO_DISK_ARRIVAL_CALLBACK_ENTRY* entry = &gDiskArrivalCallbacks[gDiskArrivalCallbackCount++];
    memset(entry, 0, sizeof(*entry));
    entry->Used = true;
    entry->Callback = callback;
    entry->Context = context;

    for (uint8_t i = 0; i < MAX_DISK_DEVICES; i++) {
        if (gDiskDevices[i])
            callback(i, gDiskDevices[i], context);
    }
    return STATUS_SUCCESS;
}

__declspec(dllexport) KSTATUS IoRegisterDiskChangeCallback(IO_DISK_CHANGE_CALLBACK callback,
                                                           void* context) {
    if (!callback)
        return STATUS_INVALID_PARAMETER;
    if (gDiskChangeCallbackCount >= MAX_DISK_CALLBACKS)
        return STATUS_NO_MEMORY;

    IO_DISK_CHANGE_CALLBACK_ENTRY* entry = &gDiskChangeCallbacks[gDiskChangeCallbackCount++];
    memset(entry, 0, sizeof(*entry));
    entry->Used = true;
    entry->Callback = callback;
    entry->Context = context;
    return STATUS_SUCCESS;
}

__declspec(dllexport) KSTATUS IoEnumerateDiskDevices(IO_DISK_ARRIVAL_CALLBACK callback,
                                                     void* context) {
    if (!callback)
        return STATUS_INVALID_PARAMETER;

    for (uint8_t i = 0; i < MAX_DISK_DEVICES; i++) {
        if (gDiskDevices[i])
            callback(i, gDiskDevices[i], context);
    }
    return STATUS_SUCCESS;
}

// Returns disk Device.
__declspec(dllexport) DEVICE_OBJECT* IoGetDiskDevice(uint8_t driveIndex) {
    if (driveIndex >= MAX_DISK_DEVICES) return (DEVICE_OBJECT*)0;
    return gDiskDevices[driveIndex];
}

// Implements the reference Disk Device operation.
__declspec(dllexport) DEVICE_OBJECT* IoReferenceDiskDevice(uint8_t driveIndex) {
    DEVICE_OBJECT* dev = IoGetDiskDevice(driveIndex);
    if (dev) IoReferenceDeviceObject(dev);
    return dev;
}

// Implements the disk Device Present operation.
__declspec(dllexport) bool IoDiskDevicePresent(uint8_t driveIndex) {
    return IoGetDiskDevice(driveIndex) != (DEVICE_OBJECT*)0;
}

// Implements the disk Sector Count operation.
__declspec(dllexport) uint64_t IoDiskSectorCount(uint8_t driveIndex) {
    DEVICE_OBJECT* dev = IoGetDiskDevice(driveIndex);
    return dev ? dev->u.Disk.SectorCount : 0;
}

// Dispatch.

__declspec(dllexport) KSTATUS IoCallDriver(DEVICE_OBJECT* device, IRP* irp) {
    if (!device || !device->DriverObject || !irp)
        return STATUS_INVALID_PARAMETER;

    IO_STACK_LOCATION* ios = IoGetCurrentIrpStackLocation(irp);
    if (!ios)
        return STATUS_INVALID_PARAMETER;

    IO_MAJOR_FUNCTION major = ios->MajorFunction;
    if (major >= IRP_MJ_MAXIMUM_FUNCTION)
        return STATUS_INVALID_PARAMETER;

    if (major != IRP_MJ_PNP &&
        (device->DeletePending || device->PnpState == DevicePnpSurpriseRemoved ||
         device->PnpState == DevicePnpRemovePending || device->PnpState == DevicePnpDeleted))
        return STATUS_OBJECT_DEAD;

    PDRIVER_DISPATCH dispatch = device->DriverObject->MajorFunction[major];
    if (!dispatch)
        return STATUS_NOT_SUPPORTED;

    ios->DeviceObject = device;
    IoReferenceDeviceObject(device);
    _InterlockedIncrement(&device->OutstandingIrps);

    KSTATUS st = dispatch(device, irp);
    if (st != STATUS_PENDING) {
        if (ios->DeviceObject == device) {
            ios->DeviceObject = (DEVICE_OBJECT*)0;
            _InterlockedDecrement(&device->OutstandingIrps);
            IoDereferenceDeviceObject(device);
        }
    }
    return st;
}

// Validates that a sector range fits the disk and IRP byte length.
static KSTATUS IoValidateDiskSectorRange(DEVICE_OBJECT* dev, uint64_t lba, uint32_t count) {
    if (!dev || count == 0) return STATUS_INVALID_PARAMETER;
    if (count > (0xffffffffu / IO_DISK_SECTOR_BYTES)) return STATUS_INVALID_PARAMETER;
    if (lba > dev->u.Disk.SectorCount) return STATUS_INVALID_PARAMETER;
    if ((uint64_t)count > dev->u.Disk.SectorCount - lba) return STATUS_INVALID_PARAMETER;
    return STATUS_SUCCESS;
}

// Submits one raw disk read IRP.
static KSTATUS IoReadDiskSectorsRaw(DEVICE_OBJECT* dev, uint64_t lba,
                                    uint32_t count, void* buffer) {
    if (!dev || !buffer || count == 0) return STATUS_INVALID_PARAMETER;

    IRP irp;
    IoBuildReadWriteIrp(&irp, IRP_MJ_READ, lba * 512ull, count * 512u, buffer);
    return IoSubmitIrpSync(dev, &irp);
}

// Submits one raw disk write IRP.
static KSTATUS IoWriteDiskSectorsRaw(DEVICE_OBJECT* dev, uint64_t lba,
                                     uint32_t count, void* buffer) {
    if (!dev || !buffer || count == 0) return STATUS_INVALID_PARAMETER;

    IRP irp;
    IoBuildReadWriteIrp(&irp, IRP_MJ_WRITE, lba * 512ull, count * 512u, buffer);
    return IoSubmitIrpSync(dev, &irp);
}

// Chooses a conservative read-ahead span from the current sequential score.
static uint32_t IoDiskChooseReadAheadSectors(uint32_t currentCount, uint32_t score) {
    uint32_t sectors = currentCount;
    if (sectors < IO_DISK_READAHEAD_MIN_SECTORS)
        sectors = IO_DISK_READAHEAD_MIN_SECTORS;
    if (score >= 8 && sectors < 128u)
        sectors = 128u;
    else if (score >= 4 && sectors < 64u)
        sectors = 64u;
    else if (score >= 2 && sectors < 32u)
        sectors = 32u;
    if (sectors > IO_DISK_READAHEAD_MAX_SECTORS)
        sectors = IO_DISK_READAHEAD_MAX_SECTORS;
    return sectors;
}

// Loads the next sequential disk range into the read-ahead buffer.
static void IoDiskMaybeReadAhead(uint8_t driveIndex, DEVICE_OBJECT* dev,
                                 uint64_t nextLba, uint32_t currentCount, uint32_t score) {
    if (driveIndex >= MAX_DISK_DEVICES || !dev || score < IO_DISK_READAHEAD_TRIGGER_SCORE)
        return;

    IO_DISK_READAHEAD_STATE* state = &gDiskReadAhead[driveIndex];
    if (!state->Buffer) {
        uintptr_t flags = KeAcquireSpinlock(&state->Lock);
        state->Stats.ReadaheadSuppressed++;
        KeReleaseSpinlock(&state->Lock, flags);
        return;
    }

    uint32_t sectors = IoDiskChooseReadAheadSectors(currentCount, score);
    uint64_t diskSectors = dev->u.Disk.SectorCount;
    if (nextLba >= diskSectors)
        return;
    if ((uint64_t)sectors > diskSectors - nextLba)
        sectors = (uint32_t)(diskSectors - nextLba);
    if (sectors == 0)
        return;

    uintptr_t flags = KeAcquireSpinlock(&state->Lock);
    bool covered = state->Valid && !state->Loading &&
                   IoDiskRangeContains(state->BaseLba, state->Count, nextLba, sectors);
    if (covered) {
        KeReleaseSpinlock(&state->Lock, flags);
        return;
    }
    if (state->Loading) {
        state->Stats.ReadaheadSuppressed++;
        KeReleaseSpinlock(&state->Lock, flags);
        return;
    }
    uint64_t generation = state->Generation;
    state->Loading = 1;
    state->Valid = 0;
    state->Count = 0;
    KeReleaseSpinlock(&state->Lock, flags);

    KSTATUS st = IoReadDiskSectorsRaw(dev, nextLba, sectors, state->Buffer);

    flags = KeAcquireSpinlock(&state->Lock);
    state->Loading = 0;
    if (KSUCCESS(st) && generation == state->Generation) {
        state->BaseLba = nextLba;
        state->Count = sectors;
        state->Valid = 1;
        state->Stats.ReadaheadStores++;
        state->Stats.ReadaheadSectorsLoaded += sectors;
        IoDiskStatMaxLocked(&state->Stats.MaxReadaheadSectors, sectors);
    } else {
        state->Valid = 0;
        state->Count = 0;
        if (!KSUCCESS(st))
            state->Stats.ReadaheadReadFailures++;
    }
    KeReleaseSpinlock(&state->Lock, flags);
}

__declspec(dllexport) KSTATUS IoReadDiskSectors(uint8_t driveIndex,
                                                uint64_t lba,
                                                uint32_t count,
                                                void* buffer) {
    if (!buffer || count == 0) return STATUS_INVALID_PARAMETER;

    DEVICE_OBJECT* dev = IoReferenceDiskDevice(driveIndex);
    if (!dev) return STATUS_NOT_FOUND;

    KSTATUS rangeStatus = IoValidateDiskSectorRange(dev, lba, count);
    if (!KSUCCESS(rangeStatus)) {
        IoDereferenceDeviceObject(dev);
        return rangeStatus;
    }

    if (IoDiskReadAheadTryHit(driveIndex, lba, count, buffer)) {
        uint32_t score = IoDiskReadAheadNoteRead(driveIndex, lba, count);
        IoDiskMaybeReadAhead(driveIndex, dev, lba + count, count, score);
        IoDereferenceDeviceObject(dev);
        return STATUS_SUCCESS;
    }

    KSTATUS st = IoReadDiskSectorsRaw(dev, lba, count, buffer);
    uint32_t score = IoDiskReadAheadNoteRead(driveIndex, lba, count);
    if (KSUCCESS(st))
        IoDiskMaybeReadAhead(driveIndex, dev, lba + count, count, score);

    IoDereferenceDeviceObject(dev);
    return st;
}

__declspec(dllexport) KSTATUS IoWriteDiskSectors(uint8_t driveIndex,
                                                 uint64_t lba,
                                                 uint32_t count,
                                                 const void* buffer) {
    if (!buffer || count == 0) return STATUS_INVALID_PARAMETER;

    DEVICE_OBJECT* dev = IoReferenceDiskDevice(driveIndex);
    if (!dev) return STATUS_NOT_FOUND;

    KSTATUS rangeStatus = IoValidateDiskSectorRange(dev, lba, count);
    if (!KSUCCESS(rangeStatus)) {
        IoDereferenceDeviceObject(dev);
        return rangeStatus;
    }

    KSTATUS st = IoWriteDiskSectorsRaw(dev, lba, count, (void*)buffer);
    if (KSUCCESS(st))
        IoDiskReadAheadNoteWrite(driveIndex, count);

    IoDereferenceDeviceObject(dev);
    return st;
}

// Queries aggregate disk I/O statistics.
__declspec(dllexport) void IoQueryDiskStats(IO_DISK_STATS* out) {
    if (!out) return;
    memset(out, 0, sizeof(*out));

    for (uint8_t i = 0; i < MAX_DISK_DEVICES; i++) {
        IO_DISK_READAHEAD_STATE* state = &gDiskReadAhead[i];
        uintptr_t flags = KeAcquireSpinlock(&state->Lock);
        if (gDiskDevices[i])
            out->RegisteredDisks++;
        out->ReadRequests += state->Stats.ReadRequests;
        out->ReadSectors += state->Stats.ReadSectors;
        out->WriteRequests += state->Stats.WriteRequests;
        out->WriteSectors += state->Stats.WriteSectors;
        if (out->MaxReadSectors < state->Stats.MaxReadSectors)
            out->MaxReadSectors = state->Stats.MaxReadSectors;
        if (out->MaxWriteSectors < state->Stats.MaxWriteSectors)
            out->MaxWriteSectors = state->Stats.MaxWriteSectors;
        out->ReadaheadHits += state->Stats.ReadaheadHits;
        out->ReadaheadMisses += state->Stats.ReadaheadMisses;
        out->ReadaheadStores += state->Stats.ReadaheadStores;
        out->ReadaheadInvalidations += state->Stats.ReadaheadInvalidations;
        out->ReadaheadBytesServed += state->Stats.ReadaheadBytesServed;
        out->ReadaheadSectorsLoaded += state->Stats.ReadaheadSectorsLoaded;
        out->ReadaheadSequentialDetections += state->Stats.ReadaheadSequentialDetections;
        out->ReadaheadSuppressed += state->Stats.ReadaheadSuppressed;
        out->ReadaheadAllocFailures += state->Stats.ReadaheadAllocFailures;
        out->ReadaheadReadFailures += state->Stats.ReadaheadReadFailures;
        if (out->MaxReadaheadSectors < state->Stats.MaxReadaheadSectors)
            out->MaxReadaheadSectors = state->Stats.MaxReadaheadSectors;
        KeReleaseSpinlock(&state->Lock, flags);
    }
}