// Kernel event primitive. This file owns manual-reset event state and waiter wakeups.

#include <stdbool.h>
#include <intrin.h>
#include <kernel/ke_event.h>
#include <kernel/ke_thread.h>
#include <kernel/ke_sched.h>
#include <kernel/ke_irql.h>
#include <kernel/ke_apc.h>
#include <kernel/ob_object.h>
#include <kernel/mm_virt.h>
#include <ntlite/syscall.h>

// Multi-wait context shared by the stack wait blocks for one wait call.
typedef struct KE_EVENT_MULTI_WAIT_CONTEXT KE_EVENT_MULTI_WAIT_CONTEXT;

typedef struct KE_EVENT_WAIT_BLOCK {
    LIST_ENTRY Link;
    KE_EVENT_MULTI_WAIT_CONTEXT* Context;
    EVENT* Event;
    uint32_t Index;
    uint8_t Armed;
} KE_EVENT_WAIT_BLOCK;

struct KE_EVENT_MULTI_WAIT_CONTEXT {
    THREAD* Thread;
    KE_EVENT_WAIT_BLOCK* Blocks;
    uint32_t Count;
    volatile long Completed;
    KSTATUS Status;
    uint32_t SelectedIndex;
};

// Implements the event Destroy operation.
static void KiEventDestroy(OBJECT* obj) { (void)obj; }

// Completes a multi-wait block while the source event lock is held.
static void KiCompleteMultiWaitBlockLocked(KE_EVENT_WAIT_BLOCK* block) {
    if (!block || !block->Armed)
        return;

    ListRemove(&block->Link);
    block->Armed = 0;

    KE_EVENT_MULTI_WAIT_CONTEXT* ctx = block->Context;
    if (!ctx)
        return;

    if (_InterlockedCompareExchange(&ctx->Completed, 1, 0) != 0)
        return;

    ctx->Status = STATUS_SUCCESS;
    ctx->SelectedIndex = block->Index;
    KeUnblockThreadFromWaitListWithStatus(ctx->Thread, STATUS_SUCCESS);
}

// Removes any wait blocks that are still linked after the wait completed.
static void KiCleanupMultiWaitBlocks(KE_EVENT_MULTI_WAIT_CONTEXT* ctx) {
    if (!ctx || !ctx->Blocks)
        return;

    for (uint32_t i = 0; i < ctx->Count; i++) {
        KE_EVENT_WAIT_BLOCK* block = &ctx->Blocks[i];
        EVENT* event = block->Event;
        if (!event)
            continue;

        uintptr_t flags = KeAcquireSpinlock(&event->Lock);
        if (block->Armed) {
            ListRemove(&block->Link);
            block->Armed = 0;
        }
        KeReleaseSpinlock(&event->Lock, flags);
    }
}

// Removes a multi-wait from every event list during abort or thread teardown.
static void KiRundownMultiWait(THREAD* thread, void* context, KSTATUS status) {
    KE_EVENT_MULTI_WAIT_CONTEXT* ctx = (KE_EVENT_MULTI_WAIT_CONTEXT*)context;
    if (!ctx)
        return;

    if (_InterlockedCompareExchange(&ctx->Completed, 1, 0) == 0) {
        ctx->Status = status;
        ctx->SelectedIndex = 0;
    }

    KiCleanupMultiWaitBlocks(ctx);

    // Abort wakes a truly blocked waiter. Thread teardown may call this after a
    // signal already made the thread ready; in that case only the list cleanup
    // above is needed.
    if (thread && thread->State == ThreadWaiting)
        KeUnblockThreadFromWaitListWithStatus(thread, status);
}

// Creates event.
__declspec(dllexport) KSTATUS KeEventCreate(EVENT_TYPE type, uint32_t initialState, EVENT **out) {
    if (!out) return STATUS_INVALID_PARAMETER;
    EVENT* e = (EVENT *)ObAllocateObject(ObjectTypeEvent, sizeof(EVENT));
    if (!e)  return STATUS_NO_MEMORY;
    e->Header.Destroy = KiEventDestroy;
    e->Lock     = (SPINLOCK)SPINLOCK_INIT;
    e->Type     = type;
    e->Signaled = initialState ? 1 : 0;
    LIST_INIT(&e->WaitList);
    LIST_INIT(&e->MultiWaitList);
    *out = e;
    return STATUS_SUCCESS;
}

// Signals event.
__declspec(dllexport) void KeSignalEvent(EVENT* event) {
    if (!event) return;
    uintptr_t flags = KeAcquireSpinlock(&event->Lock);

    if (event->Type == EventManualReset) {
        event->Signaled = 1;
        while (!LIST_IS_EMPTY(&event->WaitList)) {
            LIST_ENTRY* node = ListPopHead(&event->WaitList);
            THREAD* t = CONTAINING_RECORD(node, THREAD, SchedulerLink);
            KeSetThreadWaitList(t, (LIST_ENTRY*)0);
            KeReleaseSpinlock(&event->Lock, flags);
            KeUnblockThreadWithStatus(t, STATUS_SUCCESS);
            flags = KeAcquireSpinlock(&event->Lock);
        }
        while (!LIST_IS_EMPTY(&event->MultiWaitList)) {
            LIST_ENTRY* node = event->MultiWaitList.Fwd;
            KE_EVENT_WAIT_BLOCK* block = CONTAINING_RECORD(node, KE_EVENT_WAIT_BLOCK, Link);
            KiCompleteMultiWaitBlockLocked(block);
        }
    } else {
        while (!LIST_IS_EMPTY(&event->WaitList)) {
            event->Signaled = 0;
            LIST_ENTRY* node = ListPopHead(&event->WaitList);
            THREAD* t = CONTAINING_RECORD(node, THREAD, SchedulerLink);
            KeSetThreadWaitList(t, (LIST_ENTRY*)0);
            KeReleaseSpinlock(&event->Lock, flags);
            KeUnblockThreadWithStatus(t, STATUS_SUCCESS);
            return;
        }

        while (!LIST_IS_EMPTY(&event->MultiWaitList)) {
            LIST_ENTRY* node = event->MultiWaitList.Fwd;
            KE_EVENT_WAIT_BLOCK* block = CONTAINING_RECORD(node, KE_EVENT_WAIT_BLOCK, Link);
            KE_EVENT_MULTI_WAIT_CONTEXT* ctx = block->Context;
            if (ctx && _InterlockedCompareExchange(&ctx->Completed, 0, 0) == 0) {
                event->Signaled = 0;
                KiCompleteMultiWaitBlockLocked(block);
                KeReleaseSpinlock(&event->Lock, flags);
                return;
            }
            KiCompleteMultiWaitBlockLocked(block);
        }
        event->Signaled = 1;
    }
    KeReleaseSpinlock(&event->Lock, flags);
}

// Clears event.
__declspec(dllexport) void KeClearEvent(EVENT* event) {
    if (!event) return;
    uintptr_t flags = KeAcquireSpinlock(&event->Lock);
    event->Signaled = 0;
    KeReleaseSpinlock(&event->Lock, flags);
}

// Attempts to consume an event signal without blocking.
__declspec(dllexport) KSTATUS KeTryWaitEvent(EVENT* event) {
    if (!event) return STATUS_INVALID_PARAMETER;
    uintptr_t flags = KeAcquireSpinlock(&event->Lock);
    if (!event->Signaled) {
        KeReleaseSpinlock(&event->Lock, flags);
        return STATUS_TIMEOUT;
    }

    if (event->Type == EventAutoReset)
        event->Signaled = 0;
    KeReleaseSpinlock(&event->Lock, flags);
    return STATUS_SUCCESS;
}

// Waits for Event.
__declspec(dllexport) KSTATUS KeWaitForEvent(EVENT* event, uint8_t alertable) {
    if (!event) return STATUS_INVALID_PARAMETER;
    if (!KeCanWait()) return STATUS_INVALID_IRQL;
    if (alertable && KeAreKernelApcsDisabled()) return STATUS_INVALID_PARAMETER;
    THREAD* self = KeGetCurrentThread();

    uintptr_t flags = KeAcquireSpinlock(&event->Lock);
    if (event->Signaled) {
        if (event->Type == EventAutoReset) event->Signaled = 0;
        KeReleaseSpinlock(&event->Lock, flags);
        return STATUS_SUCCESS;
    }

    KePrepareBlockThread(self, &event->Header, alertable);
    ListInsertTail(&event->WaitList, &self->SchedulerLink);
    KeSetThreadWaitList(self, &event->WaitList);
    KeReleaseSpinlock(&event->Lock, flags);
    return KeCommitBlockThread(self, alertable);
}

// Waits for any of a set of events, returning the selected index.
__declspec(dllexport) KSTATUS KeWaitForMultipleEvents(EVENT** events, uint32_t count,
                                                     uint8_t alertable, uint64_t timeoutMs,
                                                     uint32_t* signaledIndexOut) {
    if (!events || count == 0 || count > SYS_WAIT_MANY_MAX_HANDLES)
        return STATUS_INVALID_PARAMETER;
    if (!KeCanWait()) return STATUS_INVALID_IRQL;
    if (alertable && KeAreKernelApcsDisabled()) return STATUS_INVALID_PARAMETER;

    for (uint32_t i = 0; i < count; i++) {
        if (!events[i]) return STATUS_INVALID_PARAMETER;
        KSTATUS s = KeTryWaitEvent(events[i]);
        if (KSUCCESS(s)) {
            if (signaledIndexOut) *signaledIndexOut = i;
            return STATUS_SUCCESS;
        }
    }

    if (timeoutMs == 0)
        return STATUS_TIMEOUT;

    KE_EVENT_WAIT_BLOCK stackBlocks[SYS_WAIT_MANY_MAX_HANDLES];
    KE_EVENT_MULTI_WAIT_CONTEXT ctx;
    ctx.Thread = KeGetCurrentThread();
    ctx.Blocks = stackBlocks;
    ctx.Count = count;
    ctx.Completed = 0;
    ctx.Status = STATUS_PENDING;
    ctx.SelectedIndex = 0;

    for (uint32_t i = 0; i < count; i++) {
        stackBlocks[i].Context = &ctx;
        stackBlocks[i].Event = events[i];
        stackBlocks[i].Index = i;
        stackBlocks[i].Armed = 0;
        LIST_INIT(&stackBlocks[i].Link);
    }

    KePrepareBlockThread(ctx.Thread, (OBJECT*)0, alertable);
    KeSetThreadWaitRundown(ctx.Thread, KiRundownMultiWait, &ctx);

    if (timeoutMs != SYS_WAIT_TIMEOUT_INFINITE) {
        KSTATUS ts = KeArmCurrentThreadTimeout(timeoutMs, STATUS_TIMEOUT);
        if (!KSUCCESS(ts)) {
            KeClearThreadWaitRundown(ctx.Thread);
            KeUnblockThreadFromWaitListWithStatus(ctx.Thread, ts);
            return ts;
        }
    }

    for (uint32_t i = 0; i < count; i++) {
        EVENT* event = events[i];
        uintptr_t flags = KeAcquireSpinlock(&event->Lock);
        if (_InterlockedCompareExchange(&ctx.Completed, 0, 0) == 0 && event->Signaled) {
            if (event->Type == EventAutoReset)
                event->Signaled = 0;
            ctx.Status = STATUS_SUCCESS;
            ctx.SelectedIndex = i;
            _InterlockedExchange(&ctx.Completed, 1);
            KeReleaseSpinlock(&event->Lock, flags);
            KeUnblockThreadFromWaitListWithStatus(ctx.Thread, STATUS_SUCCESS);
            break;
        }

        if (_InterlockedCompareExchange(&ctx.Completed, 0, 0) == 0) {
            stackBlocks[i].Armed = 1;
            ListInsertTail(&event->MultiWaitList, &stackBlocks[i].Link);
        }
        KeReleaseSpinlock(&event->Lock, flags);
    }

    KSTATUS waitStatus = KeCommitBlockThread(ctx.Thread, alertable);
    if (_InterlockedCompareExchange(&ctx.Completed, 1, 0) == 0)
        ctx.Status = waitStatus;

    KiCleanupMultiWaitBlocks(&ctx);
    KeClearThreadWaitRundown(ctx.Thread);

    if (ctx.Status == STATUS_SUCCESS && signaledIndexOut)
        *signaledIndexOut = ctx.SelectedIndex;
    return ctx.Status;
}