// Kernel semaphore primitive. This file owns counting wakeups and bounded resource acquisition.

#include <stdbool.h>
#include <kernel/ke_semaphore.h>
#include <kernel/ke_thread.h>
#include <kernel/ke_sched.h>
#include <kernel/ke_irql.h>
#include <kernel/ob_object.h>

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

// Implements the semaphore Create operation.
__declspec(dllexport) KSTATUS KeSemaphoreCreate(int32_t initial, int32_t max, SEMAPHORE **out) {
    if (!out || initial < 0 || max < 1 || initial > max)
        return STATUS_INVALID_PARAMETER;
    SEMAPHORE* s = (SEMAPHORE *)ObAllocateObject(ObjectTypeSemaphore, sizeof(SEMAPHORE));
    if (!s) return STATUS_NO_MEMORY;
    s->Header.Destroy = KiSemaphoreDestroy;
    s->Lock     = (SPINLOCK)SPINLOCK_INIT;
    s->Count    = initial;
    s->MaxCount = max;
    LIST_INIT(&s->WaitList);
    *out = s;
    return STATUS_SUCCESS;
}

// Waits for semaphore.
__declspec(dllexport) KSTATUS KeWaitSemaphore(SEMAPHORE* sem) {
    if (!sem) return STATUS_INVALID_PARAMETER;
    if (!KeCanWait()) return STATUS_INVALID_IRQL;
    THREAD* self = KeGetCurrentThread();

    uintptr_t flags = KeAcquireSpinlock(&sem->Lock);
    if (sem->Count > 0) {
        sem->Count--;
        KeReleaseSpinlock(&sem->Lock, flags);
        return STATUS_SUCCESS;
    }

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

// Implements the try Semaphore operation.
__declspec(dllexport) KSTATUS KeTrySemaphore(SEMAPHORE* sem) {
    if (!sem) return STATUS_INVALID_PARAMETER;
    uintptr_t flags = KeAcquireSpinlock(&sem->Lock);
    if (sem->Count > 0) {
        sem->Count--;
        KeReleaseSpinlock(&sem->Lock, flags);
        return STATUS_SUCCESS;
    }
    KeReleaseSpinlock(&sem->Lock, flags);
    return STATUS_FAILURE;
}

// Signals semaphore.
__declspec(dllexport) void KeSignalSemaphore(SEMAPHORE* sem) {
    if (!sem) return;
    uintptr_t flags = KeAcquireSpinlock(&sem->Lock);
    if (!LIST_IS_EMPTY(&sem->WaitList)) {
        // Transfer this semaphore release directly to the selected waiter. The
        // count stays unchanged; the waiter returns success without looping and
        // consuming a count later.
        LIST_ENTRY* node = ListPopHead(&sem->WaitList);
        THREAD* waiter   = CONTAINING_RECORD(node, THREAD, SchedulerLink);
        KeSetThreadWaitList(waiter, (LIST_ENTRY*)0);
        KeReleaseSpinlock(&sem->Lock, flags);
        KeUnblockThreadWithStatus(waiter, STATUS_SUCCESS);
        return;
    }
    if (sem->Count < sem->MaxCount) sem->Count++;
    KeReleaseSpinlock(&sem->Lock, flags);
}