// io_device.c - kernel driver/device dispatch and disk registry
#include <string.h>
#include <ntlite/ddi/io_driver.h>
#include <kernel/ke_irp.h>
#include <kernel/mm_virt.h>

#define MAX_DISK_DEVICES 8

static DEVICE_OBJECT* gDiskDevices[MAX_DISK_DEVICES];

__declspec(dllexport) KSTATUS IoCreateDevice(DRIVER_OBJECT* driver,
                                             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*)MmAllocatePool(total);
    if (!dev) return STATUS_NO_MEMORY;

    memset(dev, 0, total);
    dev->DriverObject = driver;
    if (extensionSize)
        dev->DeviceExtension = (void*)((uint8_t*)dev + sizeof(DEVICE_OBJECT));
    dev->AlignmentRequirement = 1;
    dev->DiskIndex = 0xFF;

    *outDevice = dev;
    return STATUS_SUCCESS;
}

__declspec(dllexport) void IoDeleteDevice(DEVICE_OBJECT* device) {
    if (!device) return;
    MmFreePool(device);
}

__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->DiskIndex = driveIndex;
    gDiskDevices[driveIndex] = device;
    return STATUS_SUCCESS;
}

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

__declspec(dllexport) bool IoDiskDevicePresent(uint8_t driveIndex) {
    return IoGetDiskDevice(driveIndex) != (DEVICE_OBJECT*)0;
}

__declspec(dllexport) uint64_t IoDiskSectorCount(uint8_t driveIndex) {
    DEVICE_OBJECT* dev = IoGetDiskDevice(driveIndex);
    return dev ? dev->SectorCount : 0;
}

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

    IO_MAJOR_FUNCTION major = (irp->Flags & IRP_FLAG_WRITE) ? IRP_MJ_WRITE : IRP_MJ_READ;
    if (major >= IRP_MJ_MAXIMUM_FUNCTION)
        return STATUS_INVALID_PARAMETER;

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

    return dispatch(device, irp);
}

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

    IRP irp;
    memset(&irp, 0, sizeof(irp));
    irp.DriveIndex = driveIndex;
    irp.Flags = 0;
    irp.Lba = lba;
    irp.SectorCount = count;
    irp.Buffer = buffer;
    irp.Status = STATUS_FAILURE;
    irp.Information = 0;

    return IoSubmitIrpSync(&irp);
}

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

    IRP irp;
    memset(&irp, 0, sizeof(irp));
    irp.DriveIndex = driveIndex;
    irp.Flags = IRP_FLAG_WRITE;
    irp.Lba = lba;
    irp.SectorCount = count;
    irp.Buffer = (void*)buffer;
    irp.Status = STATUS_FAILURE;
    irp.Information = 0;

    return IoSubmitIrpSync(&irp);
}