// 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>
#include <kernel/ke_debug.h>


#define MAX_PUBLISHED_DEVICES 64
#define MAX_DRIVER_MATCHERS 32

typedef struct IO_DRIVER_MATCHER_ENTRY {
    bool Used;
    const char* DriverName;
    IO_BUS_TYPE BusType;
    IO_DRIVER_MATCH_CALLBACK Match;
    IO_DRIVER_BIND_CALLBACK Bind;
    void* Context;
} IO_DRIVER_MATCHER_ENTRY;

static IO_PUBLISHED_DEVICE gPublishedDevices[MAX_PUBLISHED_DEVICES];
static uint32_t gPublishedDeviceCount;
static IO_DRIVER_MATCHER_ENTRY gDriverMatchers[MAX_DRIVER_MATCHERS];
static uint32_t gDriverMatcherCount;

static const char* IoBusTypeName(IO_BUS_TYPE busType) {
    switch (busType) {
        case IoBusPci: return "pci";
        case IoBusUsb: return "usb";
        default: return "unknown";
    }
}

static bool IoTryBindDevice(IO_PUBLISHED_DEVICE* device,
                            IO_DRIVER_MATCHER_ENTRY* matcher) {
    if (!device || !matcher || !matcher->Used || device->Bound)
        return false;
    if (matcher->BusType != IoBusUnknown && matcher->BusType != device->BusType)
        return false;
    if (!matcher->Match || !matcher->Bind)
        return false;
    if (!matcher->Match(device, matcher->Context))
        return false;

    KSTATUS st = matcher->Bind(device, matcher->Context);
    if (KSUCCESS(st)) {
        device->Bound = true;
        device->BoundDriverName = matcher->DriverName;
        DbgPrint("io: %s device %s bound by %s\n",
                 IoBusTypeName(device->BusType),
                 device->Name ? device->Name : "<unnamed>",
                 matcher->DriverName ? matcher->DriverName : "<unnamed>");
        return true;
    }

    DbgPrint("io: %s matcher %s accepted %s but bind failed st=0x%x\n",
             IoBusTypeName(device->BusType),
             matcher->DriverName ? matcher->DriverName : "<unnamed>",
             device->Name ? device->Name : "<unnamed>",
             st);
    return false;
}

__declspec(dllexport) KSTATUS IoRegisterDriverMatcher(const char* driverName,
                                                      IO_BUS_TYPE busType,
                                                      IO_DRIVER_MATCH_CALLBACK match,
                                                      IO_DRIVER_BIND_CALLBACK bind,
                                                      void* context) {
    if (!driverName || !match || !bind)
        return STATUS_INVALID_PARAMETER;
    if (gDriverMatcherCount >= MAX_DRIVER_MATCHERS)
        return STATUS_NO_MEMORY;

    IO_DRIVER_MATCHER_ENTRY* entry = &gDriverMatchers[gDriverMatcherCount++];
    memset(entry, 0, sizeof(*entry));
    entry->Used = true;
    entry->DriverName = driverName;
    entry->BusType = busType;
    entry->Match = match;
    entry->Bind = bind;
    entry->Context = context;

    DbgPrint("io: registered %s matcher %s\n",
             IoBusTypeName(busType), driverName);

    for (uint32_t i = 0; i < gPublishedDeviceCount; i++) {
        IoTryBindDevice(&gPublishedDevices[i], entry);
    }

    return STATUS_SUCCESS;
}

__declspec(dllexport) KSTATUS IoPublishDevice(const IO_PUBLISHED_DEVICE* device) {
    if (!device || device->BusType == IoBusUnknown)
        return STATUS_INVALID_PARAMETER;
    if (gPublishedDeviceCount >= MAX_PUBLISHED_DEVICES)
        return STATUS_NO_MEMORY;

    IO_PUBLISHED_DEVICE* stored = &gPublishedDevices[gPublishedDeviceCount++];
    memcpy(stored, device, sizeof(*stored));
    stored->Bound = false;
    stored->BoundDriverName = (const char*)0;


    for (uint32_t i = 0; i < gDriverMatcherCount; i++) {
        if (IoTryBindDevice(stored, &gDriverMatchers[i]))
            break;
    }

    return STATUS_SUCCESS;
}

__declspec(dllexport) uint32_t IoPublishedDeviceCount(void) {
    return gPublishedDeviceCount;
}

#define MAX_DISK_DEVICES 8
#define MAX_DISK_ARRIVAL_CALLBACKS 16

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

static DEVICE_OBJECT* gDiskDevices[MAX_DISK_DEVICES];
static IO_DISK_CALLBACK_ENTRY gDiskArrivalCallbacks[MAX_DISK_ARRIVAL_CALLBACKS];
static uint32_t gDiskArrivalCallbackCount;

static void IoNotifyDiskArrived(uint8_t driveIndex, DEVICE_OBJECT* device) {
    for (uint32_t i = 0; i < gDiskArrivalCallbackCount; i++) {
        IO_DISK_CALLBACK_ENTRY* entry = &gDiskArrivalCallbacks[i];
        if (entry->Used && entry->Callback)
            entry->Callback(driveIndex, device, entry->Context);
    }
}

__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;
    IoNotifyDiskArrived(driveIndex, device);
    return STATUS_SUCCESS;
}

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

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

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