// Kernel mutex primitive. This file owns exclusive ownership, wait queues, and priority-aware wakeups.

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

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

// Implements the mutex Create operation.
KSTATUS KeMutexCreate(MUTEX **out) {
    if (!out) return STATUS_INVALID_PARAMETER;
    MUTEX* m = (MUTEX *)ObAllocateObject(ObjectTypeMutex, sizeof(MUTEX));
    if (!m)  return STATUS_NO_MEMORY;
    m->Header.Destroy = KiMutexDestroy;
    m->Lock  = (SPINLOCK)SPINLOCK_INIT;
    m->Owner = (struct THREAD *)0;
    LIST_INIT(&m->WaitList);
    *out = m;
    return STATUS_SUCCESS;
}

// Implements the acquire Mutex operation.
KSTATUS KeAcquireMutex(MUTEX* mutex) {
    if (!mutex) return STATUS_INVALID_PARAMETER;
    if (!KeCanWait()) return STATUS_INVALID_IRQL;
    THREAD* self = KeGetCurrentThread();

    uintptr_t flags = KeAcquireSpinlock(&mutex->Lock);
    if (!mutex->Owner) {
        mutex->Owner = self;
        KeReleaseSpinlock(&mutex->Lock, flags);
        return STATUS_SUCCESS;
    }

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

// Implements the try Acquire Mutex operation.
KSTATUS KeTryAcquireMutex(MUTEX* mutex) {
    if (!mutex) return STATUS_INVALID_PARAMETER;
    THREAD* self = KeGetCurrentThread();
    uintptr_t flags = KeAcquireSpinlock(&mutex->Lock);
    if (!mutex->Owner) {
        mutex->Owner = self;
        KeReleaseSpinlock(&mutex->Lock, flags);
        return STATUS_SUCCESS;
    }
    KeReleaseSpinlock(&mutex->Lock, flags);
    return STATUS_FAILURE;
}

// Releases mutex.
void KeReleaseMutex(MUTEX* mutex) {
    if (!mutex) return;
    uintptr_t flags = KeAcquireSpinlock(&mutex->Lock);

    if (!LIST_IS_EMPTY(&mutex->WaitList)) {
        LIST_ENTRY* node = ListPopHead(&mutex->WaitList);
        THREAD* waiter   = CONTAINING_RECORD(node, THREAD, SchedulerLink);
        // Direct handoff: the selected waiter owns the mutex when its wait
        // returns. Do not drop Owner to NULL, or another thread can steal it.
        mutex->Owner = waiter;
        KeReleaseSpinlock(&mutex->Lock, flags);
        KeUnblockThreadWithStatus(waiter, STATUS_SUCCESS);
        return;
    }

    mutex->Owner = (struct THREAD *)0;
    KeReleaseSpinlock(&mutex->Lock, flags);
}