// I/O request packet infrastructure. This file owns IRP dispatch, queuing, cancellation, and remove-lock helpers.

#include <string.h>
#include <intrin.h>
#include <kernel/ke_irp.h>
#include <kernel/ke_event.h>
#include <kernel/ke_semaphore.h>
#include <kernel/ke_apc.h>
#include <kernel/ob_object.h>
#include <kernel/mm_virt.h>
#include <ntlite/ddi/io_driver.h>

// IRP_QUEUE operations.

__declspec(dllexport) void IrpQueueInit(IRP_QUEUE* q, struct SEMAPHORE* sem) {
    q->Lock      = (SPINLOCK)SPINLOCK_INIT;
    q->Semaphore = sem;
    LIST_INIT(&q->Head);
}

// Implements the irp Queue Enqueue operation.
__declspec(dllexport) void IrpQueueEnqueue(IRP_QUEUE* q, IRP* irp) {
    uintptr_t flags = KeAcquireSpinlock(&q->Lock);
    if (irp->Flags & IRP_FLAG_PRIORITY)
        ListInsertHead(&q->Head, &irp->Link);
    else
        ListInsertTail(&q->Head, &irp->Link);
    KeReleaseSpinlock(&q->Lock, flags);
    KeSignalSemaphore((SEMAPHORE*)q->Semaphore);
}

// Implements the irp Queue Dequeue operation.
__declspec(dllexport) IRP* IrpQueueDequeue(IRP_QUEUE* q) {
    uintptr_t flags = KeAcquireSpinlock(&q->Lock);
    LIST_ENTRY* node = ListPopHead(&q->Head);
    KeReleaseSpinlock(&q->Lock, flags);
    if (!node) return (IRP*)0;
    return CONTAINING_RECORD(node, IRP, Link);
}

__declspec(dllexport) void IrpQueueCancelAllEx(IRP_QUEUE* q,
                                                    KSTATUS status,
                                                    IRP_QUEUE_CANCEL_CALLBACK callback,
                                                    void* context) {
    if (!q) return;
    for (;;) {
        uintptr_t flags = KeAcquireSpinlock(&q->Lock);
        LIST_ENTRY* node = ListPopHead(&q->Head);
        KeReleaseSpinlock(&q->Lock, flags);
        if (!node) break;

        IRP* irp = CONTAINING_RECORD(node, IRP, Link);
        irp->Cancel = true;
        irp->Flags |= IRP_FLAG_CANCELLED;
        if (callback)
            callback(irp, context);
        IoCompleteRequest(irp, status, 0);
    }
}

// Implements the irp Queue Cancel All operation.
__declspec(dllexport) void IrpQueueCancelAll(IRP_QUEUE* q, KSTATUS status) {
    IrpQueueCancelAllEx(q, status, (IRP_QUEUE_CANCEL_CALLBACK)0, (void*)0);
}

// Remove locks.

__declspec(dllexport) KSTATUS IoInitializeRemoveLock(IO_REMOVE_LOCK* lock) {
    if (!lock) return STATUS_INVALID_PARAMETER;
    memset(lock, 0, sizeof(*lock));
    lock->Lock = (SPINLOCK)SPINLOCK_INIT;
    lock->IoCount = 1; // sentinel reference released by IoReleaseRemoveLockAndWait
    KSTATUS st = KeEventCreate(EventManualReset, 0, &lock->RemovedEvent);
    if (!KSUCCESS(st)) return st;
    return STATUS_SUCCESS;
}

// Releases remove Lock.
__declspec(dllexport) void IoDestroyRemoveLock(IO_REMOVE_LOCK* lock) {
    if (!lock) return;
    if (lock->RemovedEvent) {
        ObDereferenceObject(&lock->RemovedEvent->Header);
        lock->RemovedEvent = (EVENT*)0;
    }
}

// Implements the acquire Remove Lock operation.
__declspec(dllexport) KSTATUS IoAcquireRemoveLock(IO_REMOVE_LOCK* lock, IRP* irp) {
    (void)irp;
    if (!lock) return STATUS_INVALID_PARAMETER;
    uintptr_t flags = KeAcquireSpinlock(&lock->Lock);
    if (lock->Removing) {
        KeReleaseSpinlock(&lock->Lock, flags);
        return STATUS_OBJECT_DEAD;
    }
    _InterlockedIncrement(&lock->IoCount);
    KeReleaseSpinlock(&lock->Lock, flags);
    return STATUS_SUCCESS;
}

// Releases remove Lock.
__declspec(dllexport) void IoReleaseRemoveLock(IO_REMOVE_LOCK* lock, IRP* irp) {
    (void)irp;
    if (!lock) return;
    long left = _InterlockedDecrement(&lock->IoCount);
    if (left == 0 && lock->Removing && lock->RemovedEvent)
        KeSignalEvent(lock->RemovedEvent);
}

// Releases remove Lock And Wait.
__declspec(dllexport) void IoReleaseRemoveLockAndWait(IO_REMOVE_LOCK* lock, IRP* removeIrp) {
    (void)removeIrp;
    if (!lock) return;

    uintptr_t flags = KeAcquireSpinlock(&lock->Lock);
    lock->Removing = true;
    long left = _InterlockedDecrement(&lock->IoCount); // release sentinel
    bool done = (left == 0);
    KeReleaseSpinlock(&lock->Lock, flags);

    if (done) {
        if (lock->RemovedEvent) KeSignalEvent(lock->RemovedEvent);
        return;
    }

    if (lock->RemovedEvent)
        KeWaitForEvent(lock->RemovedEvent, 0);
}

// IRP helpers.

static size_t IoIrpSize(uint8_t stackCount) {
    if (stackCount == 0) stackCount = 1;
    return sizeof(IRP) + ((size_t)stackCount - 1u) * sizeof(IO_STACK_LOCATION);
}

// Creates irp.
__declspec(dllexport) KSTATUS IoAllocateIrp(uint8_t stackCount, IRP** outIrp) {
    if (!outIrp) return STATUS_INVALID_PARAMETER;
    if (stackCount == 0) stackCount = 1;
    IRP* irp = (IRP*)MmAllocatePoolTag(IoIrpSize(stackCount), MM_POOL_TAG('i', 'r', 'p', ' '));
    if (!irp) return STATUS_NO_MEMORY;
    IoInitializeIrp(irp, stackCount);
    *outIrp = irp;
    return STATUS_SUCCESS;
}

// Releases irp.
__declspec(dllexport) void IoFreeIrp(IRP* irp) {
    if (irp) MmFreePool(irp);
}

// Initializes irp.
__declspec(dllexport) void IoInitializeIrp(IRP* irp, uint8_t stackCount) {
    if (!irp) return;
    if (stackCount == 0) stackCount = 1;
    memset(irp, 0, IoIrpSize(stackCount));
    LIST_INIT(&irp->Link);
    irp->Status = STATUS_FAILURE;
    irp->Information = 0;
    irp->CancelLock = (SPINLOCK)SPINLOCK_INIT;
    irp->StackCount = stackCount;
    irp->CurrentLocation = 0;
}

// Returns current Irp Stack Location.
__declspec(dllexport) IO_STACK_LOCATION* IoGetCurrentIrpStackLocation(IRP* irp) {
    if (!irp || irp->CurrentLocation >= irp->StackCount) return (IO_STACK_LOCATION*)0;
    return &irp->Stack[irp->CurrentLocation];
}

// Returns next Irp Stack Location.
__declspec(dllexport) IO_STACK_LOCATION* IoGetNextIrpStackLocation(IRP* irp) {
    if (!irp || irp->CurrentLocation + 1u >= irp->StackCount) return (IO_STACK_LOCATION*)0;
    return &irp->Stack[irp->CurrentLocation + 1u];
}

// Implements the skip Current Irp Stack Location operation.
__declspec(dllexport) void IoSkipCurrentIrpStackLocation(IRP* irp) {
    if (!irp) return;
    if (irp->CurrentLocation + 1u < irp->StackCount)
        irp->CurrentLocation++;
}

// Copies current Irp Stack Location To Next.
__declspec(dllexport) void IoCopyCurrentIrpStackLocationToNext(IRP* irp) {
    IO_STACK_LOCATION* cur = IoGetCurrentIrpStackLocation(irp);
    IO_STACK_LOCATION* next = IoGetNextIrpStackLocation(irp);
    if (!cur || !next) return;
    memcpy(next, cur, sizeof(*next));
    irp->CurrentLocation++;
}

// Updates completion Routine.
__declspec(dllexport) void IoSetCompletionRoutine(IRP* irp, IRP_COMPLETION_FN fn, void* ctx) {
    if (!irp) return;
    irp->CompletionFn = fn;
    irp->CompletionCtx = ctx;
}

// Updates cancel Routine.
__declspec(dllexport) IRP_CANCEL_ROUTINE IoSetCancelRoutine(IRP* irp, IRP_CANCEL_ROUTINE routine) {
    if (!irp) return (IRP_CANCEL_ROUTINE)0;
    uintptr_t flags = KeAcquireSpinlock(&irp->CancelLock);
    IRP_CANCEL_ROUTINE old = irp->CancelRoutine;
    irp->CancelRoutine = routine;
    KeReleaseSpinlock(&irp->CancelLock, flags);
    return old;
}

// Cancels irp.
__declspec(dllexport) void IoCancelIrp(IRP* irp) {
    if (!irp) return;

    uintptr_t flags = KeAcquireSpinlock(&irp->CancelLock);
    irp->Cancel = true;
    irp->Flags |= IRP_FLAG_CANCELLED;
    IRP_CANCEL_ROUTINE routine = irp->CancelRoutine;
    irp->CancelRoutine = (IRP_CANCEL_ROUTINE)0;
    IO_STACK_LOCATION* ios = IoGetCurrentIrpStackLocation(irp);
    DEVICE_OBJECT* dev = ios ? ios->DeviceObject : (DEVICE_OBJECT*)0;
    KeReleaseSpinlock(&irp->CancelLock, flags);

    if (routine)
        routine(dev, irp);
}

__declspec(dllexport) void IoBuildReadWriteIrp(IRP* irp,
                                               IO_MAJOR_FUNCTION major,
                                               uint64_t byteOffset,
                                               uint32_t length,
                                               void* buffer) {
    if (!irp) return;
    IoInitializeIrp(irp, 1);
    irp->UserBuffer = buffer;
    IO_STACK_LOCATION* ios = IoGetCurrentIrpStackLocation(irp);
    if (!ios) return;
    ios->MajorFunction = major;
    ios->MinorFunction = 0;
    ios->Parameters.ReadWrite.ByteOffset = byteOffset;
    ios->Parameters.ReadWrite.Length = length;
}

__declspec(dllexport) void IoBuildPnpIrp(IRP* irp,
                                         IO_PNP_MINOR_FUNCTION minor) {
    if (!irp) return;
    uint8_t stackCount = irp->StackCount ? irp->StackCount : 1;
    IoInitializeIrp(irp, stackCount);
    IO_STACK_LOCATION* ios = IoGetCurrentIrpStackLocation(irp);
    if (!ios) return;
    ios->MajorFunction = IRP_MJ_PNP;
    ios->MinorFunction = (uint8_t)minor;
}

__declspec(dllexport) void IoBuildDeviceIoControlIrp(IRP* irp,
                                                     IO_MAJOR_FUNCTION major,
                                                     uint32_t ioControlCode,
                                                     void* inputBuffer,
                                                     uint32_t inputLength,
                                                     void* outputBuffer,
                                                     uint32_t outputLength) {
    if (!irp) return;
    uint8_t stackCount = irp->StackCount ? irp->StackCount : 1;
    IoInitializeIrp(irp, stackCount);
    IO_STACK_LOCATION* ios = IoGetCurrentIrpStackLocation(irp);
    if (!ios) return;
    ios->MajorFunction = major;
    ios->MinorFunction = 0;
    ios->Parameters.DeviceIoControl.IoControlCode = ioControlCode;
    ios->Parameters.DeviceIoControl.InputBuffer = inputBuffer;
    ios->Parameters.DeviceIoControl.InputLength = inputLength;
    ios->Parameters.DeviceIoControl.OutputBuffer = outputBuffer;
    ios->Parameters.DeviceIoControl.OutputLength = outputLength;
}

// Completion/submission.

__declspec(dllexport) KSTATUS IoCompleteRequest(IRP* irp, KSTATUS status, size_t information) {
    if (!irp) return STATUS_INVALID_PARAMETER;
    if (irp->Flags & IRP_FLAG_COMPLETED)
        return STATUS_ALREADY_EXISTS;

    irp->Flags |= IRP_FLAG_COMPLETED;
    irp->Status = status;
    irp->Information = information;

    IO_STACK_LOCATION* ios = IoGetCurrentIrpStackLocation(irp);
    DEVICE_OBJECT* dev = ios ? ios->DeviceObject : (DEVICE_OBJECT*)0;
    if (dev) {
        _InterlockedDecrement(&dev->OutstandingIrps);
        IoDereferenceDeviceObject(dev);
        ios->DeviceObject = (DEVICE_OBJECT*)0;
    }

    if (irp->CompletionApc)
        KeInsertQueueApc((KAPC*)irp->CompletionApc, irp, (void*)(uintptr_t)status);
    if (irp->CompletionEvent)
        KeSignalEvent((EVENT*)irp->CompletionEvent);
    if (irp->CompletionFn)
        irp->CompletionFn(irp, irp->CompletionCtx);

    return STATUS_SUCCESS;
}

// Submits irp.
__declspec(dllexport) KSTATUS IoSubmitIrp(DEVICE_OBJECT* device, IRP* irp) {
    if (!device || !irp) return STATUS_INVALID_PARAMETER;
    irp->Status = STATUS_FAILURE;
    irp->Information = 0;
    irp->Flags &= ~IRP_FLAG_COMPLETED;
    return IoCallDriver(device, irp);
}

// Submits irp Sync.
__declspec(dllexport) KSTATUS IoSubmitIrpSync(DEVICE_OBJECT* device, IRP* irp) {
    if (!device || !irp) return STATUS_INVALID_PARAMETER;

    EVENT* ev = (EVENT*)0;
    KSTATUS st = KeEventCreate(EventAutoReset, 0, &ev);
    if (!KSUCCESS(st)) return st;

    irp->CompletionEvent = ev;

    st = IoSubmitIrp(device, irp);
    if (st == STATUS_PENDING) {
        KeWaitForEvent(ev, 0);
        st = irp->Status;
    }

    ObDereferenceObject(&ev->Header);
    return st;
}