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

#include <stdbool.h>
#include <kernel/ke_list.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 NTLSTATUS_INVALID_PARAMETER;
    SEMAPHORE* s = (SEMAPHORE*)ObAllocateObject(ObjectTypeSemaphore, sizeof(SEMAPHORE));
    if (!s) return NTLSTATUS_NO_MEMORY;
    s->Header.Destroy = KiSemaphoreDestroy;
    s->Lock     = (SPINLOCK)SPINLOCK_INIT;
    s->Count    = initial;
    s->MaxCount = max;
    LIST_INIT(&s->WaitList);
    *out = s;
    return NTLSTATUS_SUCCESS;
}

// Waits for one semaphore count within the requested timeout.
__declspec(dllexport) KSTATUS KeWaitSemaphore(SEMAPHORE* sem, uint64_t timeoutMilliseconds) {
    if (!sem) return NTLSTATUS_INVALID_PARAMETER;
    if (timeoutMilliseconds != 0 && !KeCanWait()) return NTLSTATUS_INVALID_IRQL;
    THREAD* self = KeGetCurrentThread();

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

    if (timeoutMilliseconds == 0) {
        KeReleaseSpinlock(&sem->Lock, flags);
        return NTLSTATUS_TIMEOUT;
    }

    KePrepareBlockThread(self, &sem->Header, 0);
    KeListInsertTailChecked(&sem->WaitList, &self->SchedulerLink);
    KeSetThreadWaitList(self, &sem->WaitList);
    if (timeoutMilliseconds != KE_WAIT_INFINITE) {
        KSTATUS timeoutStatus = KeArmCurrentThreadTimeout(timeoutMilliseconds, NTLSTATUS_TIMEOUT);
        if (!KSUCCESS(timeoutStatus)) {
            KeListRemoveChecked(&self->SchedulerLink);
            KeSetThreadWaitList(self, (LIST_ENTRY*)0);
            KeUnblockThreadWithStatus(self, timeoutStatus);
            KeReleaseSpinlock(&sem->Lock, flags);
            return timeoutStatus;
        }
    }
    KeReleaseSpinlock(&sem->Lock, flags);
    return KeCommitBlockThread(self, 0);
}

// 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 = KeListPopHeadChecked(&sem->WaitList);
        THREAD* waiter   = CONTAINING_RECORD(node, THREAD, SchedulerLink);
        KeSetThreadWaitList(waiter, (LIST_ENTRY*)0);
        KeReleaseSpinlock(&sem->Lock, flags);
        KeUnblockThreadWithStatus(waiter, NTLSTATUS_SUCCESS);
        return;
    }
    if (sem->Count < sem->MaxCount) sem->Count++;
    KeReleaseSpinlock(&sem->Lock, flags);
}