// Kernel mutex primitive. This file owns exclusive ownership, abandonment, wait queues, and priority inheritance.

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

// Reports an invalid mutex ownership transition.
static void KiMutexBug(const char* where, MUTEX* mutex, THREAD* thread, const char* why) {
    KePanic(KE_PANIC_CODE_SCHEDULER,
            "mutex invariant failed: %s: mutex=0x%016llX thread=0x%016llX tid=%u owner=0x%016llX refs=%ld waitEmpty=%u ownerLink=0x%016llX/0x%016llX threadWaiting=0x%016llX threadOwnedEmpty=%u: %s\n",
            where,
            (uint64_t)mutex,
            (uint64_t)thread,
            thread ? thread->ThreadId : 0,
            mutex ? (uint64_t)mutex->Owner : 0,
            mutex ? mutex->Header.RefCount : 0,
            mutex ? (LIST_IS_EMPTY(&mutex->WaitList) ? 1u : 0u) : 1u,
            mutex ? (uint64_t)mutex->OwnerLink.Fwd : 0,
            mutex ? (uint64_t)mutex->OwnerLink.Bwd : 0,
            thread ? (uint64_t)thread->WaitingMutex : 0,
            thread ? (LIST_IS_EMPTY(&thread->OwnedMutexList) ? 1u : 0u) : 1u,
            why ? why : "?");
}

// Stabilizes an owner pointer before it crosses the mutex lock boundary.
static void KiMutexReferenceOwnerLocked(MUTEX* mutex, THREAD* owner) {
    if (!owner || !ObTryReferenceObject(&owner->Header))
        KiMutexBug("KiMutexReferenceOwnerLocked", mutex, owner,
                   "live mutex owner could not be referenced");
}

// Tests whether one scheduler link is currently attached to a list.
static bool KiMutexWaiterIsLinked(const THREAD* thread) {
    if (!thread) return false;
    return thread->SchedulerLink.Fwd != &thread->SchedulerLink ||
           thread->SchedulerLink.Bwd != &thread->SchedulerLink;
}

// Propagates one effective-rank change without carrying raw owners across locks.
static void KiMutexPropagatePriority(THREAD* thread) {
    int releaseThread = 0;

    for (uint32_t depth = 1; thread && depth <= KE_PRIORITY_INHERITANCE_MAX_DEPTH; depth++) {
        MUTEX* waitingMutex = (MUTEX*)0;
        bool changed = KiSchedulerRefreshEffectivePriority(thread, &waitingMutex);
        if (!changed) {
            if (releaseThread) ObDereferenceObject(&thread->Header);
            return;
        }

        KiSchedulerNotePriorityPropagation(depth);
        if (!waitingMutex) {
            if (releaseThread) ObDereferenceObject(&thread->Header);
            return;
        }

        uintptr_t flags = KeAcquireSpinlock(&waitingMutex->Lock);
        if (thread->WaitingMutex != waitingMutex ||
            thread->WaitObject != &waitingMutex->Header ||
            thread->WaitListHead != &waitingMutex->WaitList ||
            !KiMutexWaiterIsLinked(thread)) {
            KeReleaseSpinlock(&waitingMutex->Lock, flags);
            ObDereferenceObject(&waitingMutex->Header);
            if (releaseThread) ObDereferenceObject(&thread->Header);
            return;
        }

        THREAD* owner = waitingMutex->Owner;
        if (owner) KiMutexReferenceOwnerLocked(waitingMutex, owner);
        int donationChanged = KiSchedulerRefreshMutexDonation(waitingMutex);
        KeReleaseSpinlock(&waitingMutex->Lock, flags);
        ObDereferenceObject(&waitingMutex->Header);

        if (releaseThread) ObDereferenceObject(&thread->Header);
        if (!donationChanged || !owner) {
            if (owner) ObDereferenceObject(&owner->Header);
            return;
        }

        thread = owner;
        releaseThread = 1;
    }

    if (thread)
        KiMutexBug("KiMutexPropagatePriority", thread->WaitingMutex, thread,
                   "priority inheritance chain exceeded the supported depth");
}

// Initializes one mutex's ownership bookkeeping.
static void KiMutexInitialize(MUTEX* mutex) {
    mutex->Lock = (SPINLOCK)SPINLOCK_INIT;
    mutex->Owner = (THREAD*)0;
    LIST_INIT(&mutex->WaitList);
    LIST_INIT(&mutex->OwnerLink);
    mutex->DonationClass = KeSchedulingNormal;
    mutex->DonationPriority = PriorityLow;
    mutex->HasDonation = 0;
    mutex->Abandoned = 0;
}

// Destroys one unowned mutex object.
static void KiMutexDestroy(OBJECT* object) {
    MUTEX* mutex = CONTAINING_RECORD(object, MUTEX, Header);
    if (mutex->Owner || !LIST_IS_EMPTY(&mutex->WaitList) ||
        mutex->OwnerLink.Fwd != &mutex->OwnerLink ||
        mutex->OwnerLink.Bwd != &mutex->OwnerLink) {
        KiMutexBug("KiMutexDestroy", mutex, mutex->Owner,
                   "owned or waited mutex reached object destruction");
    }
}

// Assigns an unowned mutex to one thread and consumes abandonment state.
static KSTATUS KiMutexAcquireUnownedLocked(MUTEX* mutex, THREAD* thread) {
    KSTATUS status = mutex->Abandoned ? NTLSTATUS_ABANDONED : NTLSTATUS_SUCCESS;
    mutex->Abandoned = 0;
    ObReferenceObject(&mutex->Header);
    KiSchedulerAcquireMutexOwnership(thread, mutex);
    return status;
}

// Creates a mutex object.
__declspec(dllexport) KSTATUS KeMutexCreate(MUTEX** out) {
    if (!out) return NTLSTATUS_INVALID_PARAMETER;
    MUTEX* mutex = (MUTEX*)ObAllocateObject(ObjectTypeMutex, sizeof(MUTEX));
    if (!mutex) return NTLSTATUS_NO_MEMORY;
    mutex->Header.Destroy = KiMutexDestroy;
    KiMutexInitialize(mutex);
    *out = mutex;
    return NTLSTATUS_SUCCESS;
}

// Acquires a mutex within the requested timeout and donates priority while blocked.
__declspec(dllexport) KSTATUS KeAcquireMutex(MUTEX* mutex, uint64_t timeoutMilliseconds) {
    if (!mutex) return NTLSTATUS_INVALID_PARAMETER;
    if (timeoutMilliseconds != 0 && !KeCanWait()) return NTLSTATUS_INVALID_IRQL;

    THREAD* self = KeGetCurrentThread();
    if (!self) return NTLSTATUS_FAILURE;

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

    if (mutex->Owner == self) {
        KeReleaseSpinlock(&mutex->Lock, flags);
        return NTLSTATUS_DEADLOCK;
    }

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

    THREAD* owner = mutex->Owner;
    KiMutexReferenceOwnerLocked(mutex, owner);
    KSTATUS waitStatus = KiSchedulerPrepareMutexWait(self, mutex, owner);
    if (!KSUCCESS(waitStatus)) {
        KeReleaseSpinlock(&mutex->Lock, flags);
        ObDereferenceObject(&owner->Header);
        return waitStatus;
    }

    KeListInsertTailChecked(&mutex->WaitList, &self->SchedulerLink);
    if (timeoutMilliseconds != KE_WAIT_INFINITE) {
        KSTATUS timeoutStatus = KeArmCurrentThreadTimeout(timeoutMilliseconds, NTLSTATUS_TIMEOUT);
        if (!KSUCCESS(timeoutStatus)) {
            KeListRemoveChecked(&self->SchedulerLink);
            KiSchedulerCancelMutexWait(self, mutex);
            (void)KiSchedulerRefreshMutexDonation(mutex);
            KeUnblockThreadWithStatus(self, timeoutStatus);
            KeReleaseSpinlock(&mutex->Lock, flags);

            KiMutexPropagatePriority(owner);
            ObDereferenceObject(&owner->Header);
            return timeoutStatus;
        }
    }
    (void)KiSchedulerRefreshMutexDonation(mutex);
    KeReleaseSpinlock(&mutex->Lock, flags);

    KiMutexPropagatePriority(owner);
    ObDereferenceObject(&owner->Header);
    return KeCommitBlockThread(self, 0);
}

// Releases a mutex and hands ownership to its highest-ranked waiter.
__declspec(dllexport) void KeReleaseMutex(MUTEX* mutex) {
    if (!mutex) return;

    THREAD* self = KeGetCurrentThread();
    uintptr_t flags = KeAcquireSpinlock(&mutex->Lock);
    if (!self || mutex->Owner != self)
        KiMutexBug("KeReleaseMutex", mutex, self, "release by non-owner");

    THREAD* waiter = KiSchedulerSelectMutexWaiter(mutex);
    if (waiter) {
        KeListRemoveChecked(&waiter->SchedulerLink);
        KiSchedulerTransferMutexOwnership(self, waiter, mutex);
        (void)KiSchedulerRefreshMutexDonation(mutex);
        (void)KiSchedulerRefreshEffectivePriority(waiter, (MUTEX**)0);
        KeUnblockThreadWithStatus(waiter, NTLSTATUS_SUCCESS);
        KeReleaseSpinlock(&mutex->Lock, flags);

        KiMutexPropagatePriority(self);
        return;
    }

    KiSchedulerReleaseMutexOwnership(self, mutex);
    (void)KiSchedulerRefreshMutexDonation(mutex);
    KeReleaseSpinlock(&mutex->Lock, flags);

    KiMutexPropagatePriority(self);
    ObDereferenceObject(&mutex->Header);
}

// Removes one blocked mutex waiter during thread-wait rundown.
bool KiMutexAbortWait(MUTEX* mutex, THREAD* thread) {
    if (!mutex || !thread) return false;

    uintptr_t flags = KeAcquireSpinlock(&mutex->Lock);
    if (thread->WaitingMutex != mutex ||
        thread->WaitObject != &mutex->Header ||
        thread->WaitListHead != &mutex->WaitList ||
        !KiMutexWaiterIsLinked(thread)) {
        KeReleaseSpinlock(&mutex->Lock, flags);
        return false;
    }

    KeListRemoveChecked(&thread->SchedulerLink);
    KiSchedulerCancelMutexWait(thread, mutex);
    THREAD* owner = mutex->Owner;
    if (owner) KiMutexReferenceOwnerLocked(mutex, owner);
    (void)KiSchedulerRefreshMutexDonation(mutex);
    KeReleaseSpinlock(&mutex->Lock, flags);

    KiMutexPropagatePriority(owner);
    if (owner) ObDereferenceObject(&owner->Header);
    KiMutexPropagatePriority(thread);
    return true;
}

// Abandons one mutex owned by a terminating thread.
static void KiMutexAbandonOwned(MUTEX* mutex, THREAD* owner) {
    uintptr_t flags = KeAcquireSpinlock(&mutex->Lock);
    if (mutex->Owner != owner)
        KiMutexBug("KiMutexAbandonOwned", mutex, owner, "ownership list mismatch");

    THREAD* waiter = KiSchedulerSelectMutexWaiter(mutex);
    if (waiter) {
        KeListRemoveChecked(&waiter->SchedulerLink);
        KiSchedulerTransferMutexOwnership(owner, waiter, mutex);
        mutex->Abandoned = 0;
        (void)KiSchedulerRefreshMutexDonation(mutex);
        (void)KiSchedulerRefreshEffectivePriority(waiter, (MUTEX**)0);
        KeUnblockThreadWithStatus(waiter, NTLSTATUS_ABANDONED);
        KeReleaseSpinlock(&mutex->Lock, flags);

        KiMutexPropagatePriority(owner);
        return;
    }

    KiSchedulerReleaseMutexOwnership(owner, mutex);
    mutex->Abandoned = 1;
    (void)KiSchedulerRefreshMutexDonation(mutex);
    KeReleaseSpinlock(&mutex->Lock, flags);

    KiMutexPropagatePriority(owner);
    ObDereferenceObject(&mutex->Header);
}

// Releases every mutex still owned by a terminating thread.
void KiMutexRundownThread(THREAD* thread) {
    if (!thread) return;

    while (true) {
        MUTEX* mutex = KiSchedulerFirstOwnedMutex(thread);
        if (!mutex)
            return;
        KiMutexAbandonOwned(mutex, thread);
    }
}