// Kernel scheduler. This file owns run queues, dispatch decisions, preemption, and CPU-local scheduling state.

#include <ntlite/atomic.h>
#include <stddef.h>
#include <stdbool.h>
#include <intrin.h>
#include <kernel/ke_idt.h>
#include <kernel/ke_hal.h>
#include <kernel/ke_sched.h>
#include <kernel/ke_cpu.h>
#include <kernel/ke_process.h>
#include <kernel/ke_thread.h>
#include <kernel/se_security.h>
#include <kernel/ke_mutex.h>
#include <kernel/ke_spinlock.h>
#include <kernel/ke_irql.h>
#include <kernel/ke_apc.h>
#include <kernel/ke_debug.h>
#include <kernel/ke_timer.h>
#include <kernel/ke_display.h>
#include <kernel/ke_list.h>
#include <kernel/ke_work.h>
#include <kernel/mm_virt.h>

static LIST_ENTRY gNormalRunQueue[PriorityCount];
static LIST_ENTRY gRealtimeRunQueue[KE_REALTIME_PRIORITY_COUNT];
static LIST_ENTRY gTimeoutList;
static KWORK_ITEM gTimeoutWorkItem;
static SPINLOCK   gSchedLock;

static void KiTimeoutWorkRoutine(void* context);

static uint32_t gAgingCounter = 0;
static KE_SCHED_STATS gSchedStats;
static PROCESS* gIdleProcess;

// Tests whether one scheduling rank outranks another.
static bool KiRankOutranks(KE_SCHEDULING_CLASS leftClass, uint32_t leftPriority,
                          KE_SCHEDULING_CLASS rightClass, uint32_t rightPriority) {
    if (leftClass != rightClass)
        return leftClass == KeSchedulingRealtime;
    return leftPriority > rightPriority;
}

// Tests whether two scheduling ranks are equal.
static bool KiRanksEqual(KE_SCHEDULING_CLASS leftClass, uint32_t leftPriority,
                        KE_SCHEDULING_CLASS rightClass, uint32_t rightPriority) {
    return leftClass == rightClass && leftPriority == rightPriority;
}

// Compares two threads by effective scheduling rank.
static bool KiSchedulerThreadOutranks(const THREAD* left, const THREAD* right) {
    if (!left) return false;
    if (!right) return true;
    return KiRankOutranks(left->EffectiveSchedulingClass, left->EffectivePriority,
                          right->EffectiveSchedulingClass, right->EffectivePriority);
}

// Copies scheduler-owned thread state into one coherent snapshot.
void KiSchedulerSnapshotThread(THREAD* thread, KE_THREAD_SCHEDULER_SNAPSHOT* snapshot) {
    if (!thread || !snapshot) return;

    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    snapshot->State = thread->State;
    snapshot->BaseSchedulingClass = thread->BaseSchedulingClass;
    snapshot->BasePriority = thread->BasePriority;
    snapshot->EffectiveSchedulingClass = thread->EffectiveSchedulingClass;
    snapshot->EffectivePriority = thread->EffectivePriority;
    snapshot->RunningCpu = thread->RunningCpu;
    snapshot->LastCpu = thread->LastCpu;
    snapshot->QuantumRemaining = thread->QuantumRemaining;
    snapshot->QuantumTicks = thread->QuantumTicks;
    snapshot->CpuBurstTicks = thread->CpuBurstTicks;
    snapshot->LastRunTicks = thread->LastRunTicks;
    snapshot->WaitObjectType = thread->WaitObject ? (uint32_t)thread->WaitObject->Type : 0;
    snapshot->WaitStatus = thread->WaitStatus;
    snapshot->WaitingMutex = thread->WaitingMutex;
    snapshot->Queued = thread->Queued;
    snapshot->AlertableWait = thread->AlertableWait;
    KeReleaseSpinlock(&gSchedLock, flags);
}

// Selects the highest-ranked waiter while the caller holds the mutex lock.
THREAD* KiSchedulerSelectMutexWaiter(MUTEX* mutex) {
    if (!mutex) return (THREAD*)0;

    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    THREAD* best = (THREAD*)0;
    LIST_ENTRY* head = &mutex->WaitList;
    for (LIST_ENTRY* node = head->Fwd; node != head; node = node->Fwd) {
        THREAD* waiter = CONTAINING_RECORD(node, THREAD, SchedulerLink);
        if (!best || KiSchedulerThreadOutranks(waiter, best))
            best = waiter;
    }
    KeReleaseSpinlock(&gSchedLock, flags);
    return best;
}

// Refreshes one mutex donation while the caller holds the mutex lock.
bool KiSchedulerRefreshMutexDonation(MUTEX* mutex) {
    if (!mutex) return false;

    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    THREAD* highest = (THREAD*)0;
    LIST_ENTRY* head = &mutex->WaitList;
    for (LIST_ENTRY* node = head->Fwd; node != head; node = node->Fwd) {
        THREAD* waiter = CONTAINING_RECORD(node, THREAD, SchedulerLink);
        if (!highest || KiSchedulerThreadOutranks(waiter, highest))
            highest = waiter;
    }

    KE_SCHEDULING_CLASS oldClass = mutex->DonationClass;
    uint32_t oldPriority = mutex->DonationPriority;
    uint8_t oldHasDonation = mutex->HasDonation;
    if (highest) {
        mutex->DonationClass = highest->EffectiveSchedulingClass;
        mutex->DonationPriority = highest->EffectivePriority;
        mutex->HasDonation = 1;
    } else {
        mutex->DonationClass = KeSchedulingNormal;
        mutex->DonationPriority = PriorityLow;
        mutex->HasDonation = 0;
    }

    bool changed = oldHasDonation != mutex->HasDonation ||
                  oldClass != mutex->DonationClass ||
                  oldPriority != mutex->DonationPriority;
    KeReleaseSpinlock(&gSchedLock, flags);
    return changed;
}

// Implements the default Thread Affinity Mask operation.
static uint64_t KiDefaultThreadAffinityMask(void) {
    // Normal kernel, user, and driver-created threads are boot-processor-affine by
    // default. Code that is explicitly audited for AP execution can widen its
    // affinity before the thread enters the ready queue.
    return 1ULL;
}

// Implements the thread Is Killed Locked operation.
static bool KiThreadIsKilledLocked(const THREAD* thread) {
    return thread && !thread->IsIdleThread && thread->Owner &&
           KeIsProcessTerminating(thread->Owner) &&
           !KeThreadTerminationDeferred(thread);
}

// Implements the thread Can Run On Cpu Locked operation.
static bool KiThreadCanRunOnCpuLocked(const THREAD* thread, const KCPU* cpu) {
    if (!thread || !cpu) return false;
    if (thread->IsIdleThread) return thread == cpu->IdleThread;
    if (cpu->CpuIndex >= 64) return false;
    return (thread->AffinityMask & (1ULL << cpu->CpuIndex)) != 0;
}

// Implements the cpu Is Idle Locked operation.
static bool KiCpuIsIdleLocked(const KCPU* cpu) {
    if (!cpu || cpu->State != KeCpuOnline)
        return false;
    THREAD* cur = cpu->CurrentThread;
    return (!cur || cur == cpu->IdleThread || cur->State != ThreadRunning);
}

// Implements the any Idle Primary For Thread Locked operation.
static bool KiAnyIdlePrimaryForThreadLocked(const THREAD* thread) {
    if (!thread || thread->IsIdleThread)
        return false;

    for (uint32_t i = 0; i < gCpuCount; i++) {
        KCPU* cpu = &gCpus[i];
        if (cpu->State != KeCpuOnline)
            continue;
        if (!cpu->IsPrimaryThread)
            continue;
        if (!KiCpuIsIdleLocked(cpu))
            continue;
        if (KiThreadCanRunOnCpuLocked(thread, cpu))
            return true;
    }
    return false;
}

// Checks whether defer To Idle Primary Locked.
static bool KiShouldDeferToIdlePrimaryLocked(KCPU* cpu, const THREAD* thread) {
    if (!cpu || !thread || thread->IsIdleThread)
        return false;
    if (cpu->IsPrimaryThread || cpu->SmtThreadCount <= 1)
        return false;

    // Do not fill a secondary SMT sibling while any primary hardware thread
    // that can run this thread is idle. KiPokeRunnableCpusLocked targets those
    // primary CPUs first, so this CPU can safely remain idle and wait for more
    // work or for physical cores to fill.
    return KiAnyIdlePrimaryForThreadLocked(thread);
}

// Implements the request Reschedule Cpu Locked operation.
static void KiRequestRescheduleCpuLocked(KCPU* target) {
    if (!target || target->State != KeCpuOnline)
        return;

    NtlAtomicExchange32(&target->ReschedulePending, 1);

    KCPU* self = KeGetCurrentCpu();
    if (self && target->CpuIndex != self->CpuIndex) {
        KeSendProcessorInterrupt(target->HardwareId, KE_VECTOR_RESCHEDULE_IPI);
    }
}

// Implements the poke Runnable Cpus Locked operation.
static void KiPokeRunnableCpusLocked(THREAD* thread) {
    if (!thread || thread->IsIdleThread)
        return;

    int haveIdlePrimary = KiAnyIdlePrimaryForThreadLocked(thread);
    int requestedIdle = 0;

    for (uint32_t i = 0; i < gCpuCount; i++) {
        KCPU* cpu = &gCpus[i];
        if (cpu->State != KeCpuOnline)
            continue;
        if (!KiThreadCanRunOnCpuLocked(thread, cpu))
            continue;

        THREAD* cur = cpu->CurrentThread;
        if (!cur || cur == cpu->IdleThread) {
            if (haveIdlePrimary && !cpu->IsPrimaryThread)
                continue;
            if (cpu->IsPrimaryThread)
                gSchedStats.SmtPrimaryPokes++;
            KiRequestRescheduleCpuLocked(cpu);
            requestedIdle = 1;
        } else if (cur->State == ThreadRunning && KiSchedulerThreadOutranks(thread, cur)) {
            gSchedStats.PriorityPreemptions++;
            if (thread->EffectiveSchedulingClass == KeSchedulingRealtime)
                gSchedStats.RealtimePreemptions++;
            KiRequestRescheduleCpuLocked(cpu);
        }
    }

    // If no preferred idle primary was actually targetable, allow an idle SMT
    // sibling to take the work. This covers explicit affinity to a secondary
    // logical CPU and partially-online systems.
    if (!requestedIdle && haveIdlePrimary == 0) {
        for (uint32_t i = 0; i < gCpuCount; i++) {
            KCPU* cpu = &gCpus[i];
            if (cpu->State != KeCpuOnline)
                continue;
            if (!KiThreadCanRunOnCpuLocked(thread, cpu))
                continue;
            THREAD* cur = cpu->CurrentThread;
            if (!cur || cur == cpu->IdleThread) {
                KiRequestRescheduleCpuLocked(cpu);
                break;
            }
        }
    }
}

// Implements the base Quantum For Priority operation.
static uint32_t KiBaseQuantumForPriority(uint32_t priority) {
    switch ((THREAD_PRIORITY)priority) {
    case PriorityHigh:   return SCHED_HIGH_QUANTUM_TICKS;
    case PriorityNormal: return SCHED_NORMAL_QUANTUM_TICKS;
    case PriorityLow:    return SCHED_LOW_QUANTUM_TICKS;
    default:             return SCHED_NORMAL_QUANTUM_TICKS;
    }
}

// Implements the compute Quantum Locked operation.
static uint32_t KiComputeQuantumLocked(THREAD* thread) {
    if (!thread || thread == KeGetCurrentCpu()->IdleThread)
        return 1;

    if (thread->EffectiveSchedulingClass == KeSchedulingRealtime)
        return SCHED_REALTIME_QUANTUM_TICKS;

    uint32_t quantum = KiBaseQuantumForPriority(thread->EffectivePriority);

    // Low-priority CPU-bound threads get longer batches once they have proven
    // they mostly burn CPU. Higher-priority work still preempts them as soon
    // as it becomes ready, so this reduces switch overhead without hurting
    // interactive wake latency.
    if (thread->EffectivePriority == PriorityLow && thread->CpuBurstTicks >=
        (SCHED_DEMOTE_AFTER_BURST_TICKS * 2u)) {
        quantum += 4;
    }

    return quantum ? quantum : 1;
}

// Implements the refill Quantum Locked operation.
static void KiRefillQuantumLocked(THREAD* thread) {
    if (!thread) return;
    thread->QuantumTicks = KiComputeQuantumLocked(thread);
    thread->QuantumRemaining = thread->QuantumTicks;
}

// Returns current Thread.
THREAD* KeGetCurrentThread(void)      { return KeGetCurrentCpu()->CurrentThread; }
// Updates current Thread.
void    KeSetCurrentThread(THREAD* t) { KeGetCurrentCpu()->CurrentThread = t; }

// Returns scheduler Stats.
void KeGetSchedulerStats(KE_SCHED_STATS* out) {
    if (!out) return;
    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    *out = gSchedStats;
    KeReleaseSpinlock(&gSchedLock, flags);
}

// Implements the arm Post Switch Unlock operation.
static void KiArmPostSwitchUnlock(THREAD* thread, uintptr_t cookie) {
    // The scheduler lock is intentionally held across KeSwitchContext(). The
    // unlock cookie belongs to the thread continuation that will resume after
    // the switch, not to the physical CPU. Another dispatch may resume this
    // thread from an unrelated interrupt, so using a KCPU slot here would
    // restore the outgoing CPU context's IF/IRQL into the incoming thread.
    if (thread)
        thread->SwitchLockCookie = cookie;
}

// Implements the post Switch Unlock Cookie operation.
static uintptr_t KiPostSwitchUnlockCookie(THREAD* thread) {
    return thread ? thread->SwitchLockCookie : 0x202ULL;
}

// Implements the scheduler Bug operation.
static void KiSchedulerBug(const char* where, THREAD* thread, const char* why) {
    KePanic(KE_PANIC_CODE_SCHEDULER,
            "scheduler invariant failed: %s: thread=0x%016llX tid=%u state=%u queued=%u runningCpu=%u idle=%u: %s\n",
            where,
            (uint64_t)thread,
            thread ? thread->ThreadId : 0,
            thread ? (unsigned)thread->State : 0,
            thread ? (unsigned)thread->Queued : 0,
            thread ? (unsigned)thread->RunningCpu : 0,
            thread ? (unsigned)thread->IsIdleThread : 0,
            why ? why : "?");
}

// Implements the assert Ready Queue Candidate Locked operation.
static void KiAssertReadyQueueCandidateLocked(const char* where, THREAD* thread) {
    if (!thread)
        KiSchedulerBug(where, thread, "null thread");
    if (thread->IsIdleThread)
        KiSchedulerBug(where, thread, "idle thread cannot enter ready queue");
    if (thread->State == ThreadDead)
        KiSchedulerBug(where, thread, "dead thread cannot enter ready queue");
    if (thread->Queued)
        KiSchedulerBug(where, thread, "thread already in ready queue");
    if (thread->RunningCpu != KE_INVALID_CPU)
        KiSchedulerBug(where, thread, "thread still reserved/running on a CPU");
}

// Implements the assert Picked Thread Locked operation.
static void KiAssertPickedThreadLocked(const char* where, THREAD* thread) {
    if (!thread)
        KiSchedulerBug(where, thread, "null picked thread");
    if (thread->IsIdleThread)
        return;
    if (thread->Queued)
        KiSchedulerBug(where, thread, "picked thread still marked queued");
    if (thread->RunningCpu != KE_INVALID_CPU)
        KiSchedulerBug(where, thread, "picked thread already running on a CPU");
    if (thread->State != ThreadReady)
        KiSchedulerBug(where, thread, "picked thread was not ready");
}

// Inserts one thread into the queue selected by its effective rank.
static void KiInsertReadyLocked(THREAD* thread) {
    KiAssertReadyQueueCandidateLocked("KiInsertReadyLocked", thread);

    LIST_ENTRY* queue;
    if (thread->EffectiveSchedulingClass == KeSchedulingRealtime) {
        uint32_t priority = thread->EffectivePriority;
        if (priority > KE_REALTIME_PRIORITY_MAX)
            priority = KE_REALTIME_PRIORITY_MAX;
        queue = &gRealtimeRunQueue[priority];
    } else {
        uint32_t priority = thread->EffectivePriority;
        if (priority >= PriorityCount)
            priority = PriorityCount - 1;
        queue = &gNormalRunQueue[priority];
    }

    thread->State = ThreadReady;
    thread->Queued = 1;
    KeListInsertTailChecked(queue, &thread->SchedulerLink);
}

// Tests whether any ready thread outranks one running thread on a CPU.
static bool KiAnyRunnableOutrankingThreadForCpuLocked(KCPU* cpu, const THREAD* current) {
    if (!cpu || !current) return false;

    for (int p = (int)KE_REALTIME_PRIORITY_MAX; p >= 0; p--) {
        LIST_ENTRY* head = &gRealtimeRunQueue[p];
        for (LIST_ENTRY* node = head->Fwd; node != head; node = node->Fwd) {
            THREAD* thread = CONTAINING_RECORD(node, THREAD, SchedulerLink);
            if (!KiThreadCanRunOnCpuLocked(thread, cpu))
                continue;
            if (KiSchedulerThreadOutranks(thread, current))
                return true;
        }
    }

    for (int p = PriorityCount - 1; p >= 0; p--) {
        LIST_ENTRY* head = &gNormalRunQueue[p];
        for (LIST_ENTRY* node = head->Fwd; node != head; node = node->Fwd) {
            THREAD* thread = CONTAINING_RECORD(node, THREAD, SchedulerLink);
            if (!KiThreadCanRunOnCpuLocked(thread, cpu))
                continue;
            if (KiSchedulerThreadOutranks(thread, current))
                return true;
        }
    }

    return false;
}

// Tests whether a real-time peer at the current or higher priority is ready.
static bool KiAnyRealtimePeerForCpuLocked(KCPU* cpu, uint32_t minimumPriority) {
    if (!cpu) return false;
    if (minimumPriority > KE_REALTIME_PRIORITY_MAX)
        minimumPriority = KE_REALTIME_PRIORITY_MAX;

    for (int p = (int)KE_REALTIME_PRIORITY_MAX; p >= (int)minimumPriority; p--) {
        LIST_ENTRY* head = &gRealtimeRunQueue[p];
        for (LIST_ENTRY* node = head->Fwd; node != head; node = node->Fwd) {
            THREAD* thread = CONTAINING_RECORD(node, THREAD, SchedulerLink);
            if (KiThreadCanRunOnCpuLocked(thread, cpu))
                return true;
        }
    }
    return false;
}

// Computes one thread's effective rank from its base rank and mutex donations.
static void KiComputeEffectiveRankLocked(THREAD* thread,
                                         KE_SCHEDULING_CLASS* schedulingClassOut,
                                         uint32_t* priorityOut) {
    KE_SCHEDULING_CLASS schedulingClass = thread->BaseSchedulingClass;
    uint32_t priority = thread->BasePriority;

    LIST_ENTRY* head = &thread->OwnedMutexList;
    for (LIST_ENTRY* node = head->Fwd; node != head; node = node->Fwd) {
        MUTEX* mutex = CONTAINING_RECORD(node, MUTEX, OwnerLink);
        if (!mutex->HasDonation)
            continue;
        if (KiRankOutranks(mutex->DonationClass, mutex->DonationPriority,
                           schedulingClass, priority)) {
            schedulingClass = mutex->DonationClass;
            priority = mutex->DonationPriority;
        }
    }

    *schedulingClassOut = schedulingClass;
    *priorityOut = priority;
}

// Applies one effective-rank change and preserves ready-queue ordering.
static bool KiApplyEffectiveRankLocked(THREAD* thread,
                                      KE_SCHEDULING_CLASS schedulingClass,
                                      uint32_t priority) {
    KE_SCHEDULING_CLASS oldClass = thread->EffectiveSchedulingClass;
    uint32_t oldPriority = thread->EffectivePriority;
    if (KiRanksEqual(oldClass, oldPriority, schedulingClass, priority))
        return false;

    int wasQueued = thread->Queued != 0;
    if (wasQueued) {
        KeListRemoveChecked(&thread->SchedulerLink);
        thread->Queued = 0;
    }

    thread->EffectiveSchedulingClass = schedulingClass;
    thread->EffectivePriority = priority;
    thread->QuantumRemaining = 0;

    if (wasQueued) {
        KiInsertReadyLocked(thread);
        KiPokeRunnableCpusLocked(thread);
    } else if (thread->State == ThreadRunning && thread->RunningCpu < gCpuCount) {
        KCPU* cpu = &gCpus[thread->RunningCpu];
        if (KiAnyRunnableOutrankingThreadForCpuLocked(cpu, thread))
            KiRequestRescheduleCpuLocked(cpu);
    }

    return true;
}

// Recomputes effective rank from current base rank and owned mutexes.
bool KiSchedulerRefreshEffectivePriority(THREAD* thread, MUTEX** waitingMutexOut) {
    if (!thread) return false;

    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    KE_SCHEDULING_CLASS oldClass = thread->EffectiveSchedulingClass;
    uint32_t oldPriority = thread->EffectivePriority;
    KE_SCHEDULING_CLASS schedulingClass;
    uint32_t priority;
    KiComputeEffectiveRankLocked(thread, &schedulingClass, &priority);
    bool changed = KiApplyEffectiveRankLocked(thread, schedulingClass, priority);
    if (changed) {
        if (KiRankOutranks(schedulingClass, priority, oldClass, oldPriority))
            gSchedStats.PriorityInheritanceBoosts++;
        else
            gSchedStats.PriorityInheritanceDrops++;
    }
    if (waitingMutexOut) {
        *waitingMutexOut = (MUTEX*)0;
        if (changed && thread->WaitingMutex) {
            if (!ObTryReferenceObject(&thread->WaitingMutex->Header))
                KiSchedulerBug("KiSchedulerRefreshEffectivePriority", thread,
                               "waiting mutex could not be referenced");
            *waitingMutexOut = thread->WaitingMutex;
        }
    }
    KeReleaseSpinlock(&gSchedLock, flags);
    return changed;
}

// Publishes one wait edge atomically so concurrent acquisitions cannot create a cycle.
KSTATUS KiSchedulerPrepareMutexWait(THREAD* thread, MUTEX* mutex, THREAD* owner) {
    if (!thread || !mutex || !owner) return NTLSTATUS_INVALID_PARAMETER;

    KCPU* cpu = KeGetCurrentCpu();
    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);

    THREAD* current = owner;
    uint32_t depth = 0;
    while (current && depth < KE_PRIORITY_INHERITANCE_MAX_DEPTH) {
        if (current == thread) {
            KeReleaseSpinlock(&gSchedLock, flags);
            return NTLSTATUS_DEADLOCK;
        }

        MUTEX* waitingMutex = current->WaitingMutex;
        current = waitingMutex ? waitingMutex->Owner : (THREAD*)0;
        depth++;
    }

    if (current) {
        KeReleaseSpinlock(&gSchedLock, flags);
        return NTLSTATUS_DEADLOCK;
    }

    if (thread->IsIdleThread)
        KiSchedulerBug("KiSchedulerPrepareMutexWait", thread, "idle thread cannot block");
    if (thread->Queued)
        KiSchedulerBug("KiSchedulerPrepareMutexWait", thread, "queued thread cannot block");
    if (!cpu || thread != cpu->CurrentThread)
        KiSchedulerBug("KiSchedulerPrepareMutexWait", thread, "only the current thread can prepare a mutex wait");
    if (thread->RunningCpu != cpu->CpuIndex)
        KiSchedulerBug("KiSchedulerPrepareMutexWait", thread, "thread running CPU does not match current CPU");
    if (thread->State != ThreadRunning || thread->WaitingMutex || thread->WaitObject ||
        thread->WaitListHead || thread->SchedulerLink.Fwd != &thread->SchedulerLink ||
        thread->SchedulerLink.Bwd != &thread->SchedulerLink || thread->TimeoutArmed ||
        thread->TimeoutLink.Fwd != &thread->TimeoutLink ||
        thread->TimeoutLink.Bwd != &thread->TimeoutLink) {
        KiSchedulerBug("KiSchedulerPrepareMutexWait", thread, "thread already has wait state");
    }

    thread->State = ThreadWaiting;
    thread->WaitObject = &mutex->Header;
    thread->WaitStatus = NTLSTATUS_PENDING;
    thread->WaitListHead = &mutex->WaitList;
    thread->AlertableWait = 0;
    thread->WaitRundownRoutine = (KE_WAIT_RUNDOWN_ROUTINE)0;
    thread->WaitRundownContext = (void*)0;
    thread->WaitingMutex = mutex;
    KeReleaseSpinlock(&gSchedLock, flags);
    return NTLSTATUS_SUCCESS;
}

// Removes one published wait edge after its wait-list node is detached.
void KiSchedulerCancelMutexWait(THREAD* thread, MUTEX* mutex) {
    if (!thread || !mutex) return;

    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    if (thread->WaitingMutex != mutex || thread->WaitObject != &mutex->Header ||
        thread->WaitListHead != &mutex->WaitList ||
        thread->SchedulerLink.Fwd != &thread->SchedulerLink ||
        thread->SchedulerLink.Bwd != &thread->SchedulerLink) {
        KiSchedulerBug("KiSchedulerCancelMutexWait", thread, "mutex wait state is inconsistent");
    }

    thread->WaitingMutex = (MUTEX*)0;
    thread->WaitListHead = (LIST_ENTRY*)0;
    KeReleaseSpinlock(&gSchedLock, flags);
}

// Assigns one unowned mutex while the caller holds its lock.
void KiSchedulerAcquireMutexOwnership(THREAD* thread, MUTEX* mutex) {
    if (!thread || !mutex) return;

    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    if (mutex->Owner ||
        mutex->OwnerLink.Fwd != &mutex->OwnerLink ||
        mutex->OwnerLink.Bwd != &mutex->OwnerLink) {
        KiSchedulerBug("KiSchedulerAcquireMutexOwnership", thread,
                       "mutex already belongs to an owner");
    }

    mutex->Owner = thread;
    KeListInsertTailChecked(&thread->OwnedMutexList, &mutex->OwnerLink);
    KeReleaseSpinlock(&gSchedLock, flags);
}

// Transfers one mutex directly to a blocked waiter.
void KiSchedulerTransferMutexOwnership(THREAD* oldOwner, THREAD* newOwner,
                                       MUTEX* mutex) {
    if (!oldOwner || !newOwner || !mutex) return;

    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    if (mutex->Owner != oldOwner ||
        mutex->OwnerLink.Fwd == &mutex->OwnerLink ||
        mutex->OwnerLink.Bwd == &mutex->OwnerLink) {
        KiSchedulerBug("KiSchedulerTransferMutexOwnership", oldOwner,
                       "mutex ownership graph is inconsistent");
    }
    if (newOwner->WaitingMutex != mutex ||
        newOwner->State != ThreadWaiting ||
        newOwner->WaitObject != &mutex->Header ||
        newOwner->WaitListHead != &mutex->WaitList ||
        newOwner->SchedulerLink.Fwd != &newOwner->SchedulerLink ||
        newOwner->SchedulerLink.Bwd != &newOwner->SchedulerLink) {
        KiSchedulerBug("KiSchedulerTransferMutexOwnership", newOwner,
                       "new owner was not in the mutex handoff state");
    }

    KeListRemoveChecked(&mutex->OwnerLink);
    newOwner->WaitingMutex = (MUTEX*)0;
    newOwner->WaitListHead = (LIST_ENTRY*)0;
    mutex->Owner = newOwner;
    KeListInsertTailChecked(&newOwner->OwnedMutexList, &mutex->OwnerLink);
    KeReleaseSpinlock(&gSchedLock, flags);
}

// Removes one mutex from its current owner.
void KiSchedulerReleaseMutexOwnership(THREAD* thread, MUTEX* mutex) {
    if (!thread || !mutex) return;

    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    if (mutex->Owner != thread ||
        mutex->OwnerLink.Fwd == &mutex->OwnerLink ||
        mutex->OwnerLink.Bwd == &mutex->OwnerLink) {
        KiSchedulerBug("KiSchedulerReleaseMutexOwnership", thread,
                       "mutex was not linked to the stated owner");
    }

    KeListRemoveChecked(&mutex->OwnerLink);
    mutex->Owner = (THREAD*)0;
    KeReleaseSpinlock(&gSchedLock, flags);
}

// Returns one owned mutex while the ownership list is stable.
MUTEX* KiSchedulerFirstOwnedMutex(THREAD* thread) {
    if (!thread) return (MUTEX*)0;
    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    MUTEX* mutex = LIST_IS_EMPTY(&thread->OwnedMutexList)
        ? (MUTEX*)0
        : CONTAINING_RECORD(thread->OwnedMutexList.Fwd, MUTEX, OwnerLink);
    KeReleaseSpinlock(&gSchedLock, flags);
    return mutex;
}

// Records one transitive priority-inheritance propagation step.
void KiSchedulerNotePriorityPropagation(uint32_t depth) {
    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    gSchedStats.PriorityInheritancePropagations++;
    if (depth > gSchedStats.PriorityInheritanceMaxDepth)
        gSchedStats.PriorityInheritanceMaxDepth = depth;
    KeReleaseSpinlock(&gSchedLock, flags);
}

// Changes one normal thread's dynamic base priority.
KSTATUS KiSchedulerSetNormalPriority(THREAD* thread, THREAD_PRIORITY priority) {
    if (!thread || priority >= PriorityCount)
        return NTLSTATUS_INVALID_PARAMETER;

    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    if (thread->BaseSchedulingClass != KeSchedulingNormal) {
        KeReleaseSpinlock(&gSchedLock, flags);
        return NTLSTATUS_ACCESS_DENIED;
    }

    thread->BasePriority = (uint32_t)priority;
    KE_SCHEDULING_CLASS schedulingClass;
    uint32_t effectivePriority;
    KiComputeEffectiveRankLocked(thread, &schedulingClass, &effectivePriority);
    (void)KiApplyEffectiveRankLocked(thread, schedulingClass, effectivePriority);
    KeReleaseSpinlock(&gSchedLock, flags);
    return NTLSTATUS_SUCCESS;
}

static bool KiAnyRunnableForCpuLocked(KCPU* cpu);

// Implements the any Runnable For Current Cpu operation.
static int KiAnyRunnableForCurrentCpu(void) {
    KCPU* cpu = KeGetCurrentCpu();
    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    int any = KiAnyRunnableForCpuLocked(cpu);
    KeReleaseSpinlock(&gSchedLock, flags);
    return any;
}

// Implements the idle Loop operation.
static void KiIdleLoop(void) {
    while (true) {
        KCPU* cpu = KeGetCurrentCpu();
        cpu->SchedulerTicks++;

        // Secondary-processor idle threads sleep until an interrupt or reschedule IPI wakes
        // the CPU. Ready paths set ReschedulePending and target that processor;
        // KiExitInterrupt consumes the request after EOI.
        if (KiAnyRunnableForCurrentCpu()) {
            KeYield();
            continue;
        }

        // HLT only wakes for maskable IPIs when IF=1. A CPU that reaches
        // the idle thread through an unusual switch path must not sleep with
        // interrupts disabled, otherwise reschedule/TLB shootdown IPIs can
        // be delayed forever.
        KeWaitForInterrupt();
    }
}

// Initializes scheduler Global.
void KiInitializeSchedulerGlobal(void) {
    for (int i = 0; i < PriorityCount; i++) LIST_INIT(&gNormalRunQueue[i]);
    for (uint32_t i = 0; i < KE_REALTIME_PRIORITY_COUNT; i++) LIST_INIT(&gRealtimeRunQueue[i]);
    LIST_INIT(&gTimeoutList);
    KeInitializeWorkItem(&gTimeoutWorkItem, KiTimeoutWorkRoutine, (void*)0);
    gSchedLock = (SPINLOCK)SPINLOCK_INIT;

    KSTATUS s = KeCreateProcess("idle", SeGetRootSecurityContext(), &gIdleProcess);
    if (!KSUCCESS(s))
        KePanic(KE_PANIC_CODE_INIT, "scheduler: failed to create idle process, status=%u\n", s);
    gIdleProcess->Flags |= KE_PROCESS_FLAG_IDLE | KE_PROCESS_FLAG_TERMINATION_PROTECTED;

    PROCESS* systemProc;
    s = KeCreateSystemProcess(&systemProc);
    if (!KSUCCESS(s))
        KePanic(KE_PANIC_CODE_INIT, "scheduler: failed to create system process, status=%u\n", s);
    (void)systemProc;
}

// Creates idle Thread For Cpu.
KSTATUS KiCreateIdleThreadForCpu(KCPU* cpu) {
    if (!cpu || !gIdleProcess)
        return NTLSTATUS_INVALID_PARAMETER;
    if (cpu->IdleThread)
        return NTLSTATUS_SUCCESS;

    THREAD* idle = (THREAD*)ObAllocateObject(ObjectTypeThread, sizeof(THREAD));
    if (!idle)
        return NTLSTATUS_NO_MEMORY;

    idle->Header.Destroy = (void*)0;
    KSTATUS idStatus = KeAllocateThreadId(&idle->ThreadId);
    if (!KSUCCESS(idStatus)) {
        ObDereferenceObject(&idle->Header);
        return idStatus;
    }
    idle->State                    = ThreadRunning;
    idle->BaseSchedulingClass      = KeSchedulingNormal;
    idle->BasePriority             = PriorityLow;
    idle->EffectiveSchedulingClass = KeSchedulingNormal;
    idle->EffectivePriority        = PriorityLow;
    idle->QuantaConsumed           = 0;
    idle->QuantumTicks     = 1;
    idle->QuantumRemaining = 1;
    idle->CpuBurstTicks    = 0;
    idle->LastRunTicks     = 0;
    idle->RunningCpu       = cpu->CpuIndex;
    idle->SuspendCount     = 0;
    LIST_INIT(&idle->OwnedMutexList);
    idle->WaitingMutex     = (MUTEX*)0;
    idle->Queued           = 0;
    idle->IsIdleThread     = 1;
    idle->SchedulerFlags   = 0;
    idle->LastCpu          = cpu->CpuIndex;
    idle->AffinityMask     = (cpu->CpuIndex < 64) ? (1ULL << cpu->CpuIndex) : 0ULL;
    idle->Owner           = gIdleProcess;
    idle->WaitObject      = (OBJECT*)0;
    idle->WaitStatus      = NTLSTATUS_SUCCESS;
    idle->WaitListHead    = (LIST_ENTRY*)0;
    idle->AlertableWait   = 0;
    LIST_INIT(&idle->TimeoutLink);
    idle->TimeoutArmed    = 0;
    idle->TimeoutTick     = 0;
    idle->TimeoutStatus   = NTLSTATUS_SUCCESS;
    idle->UserStackBase   = (void*)0;
    idle->Teb             = (void*)0;
    idle->TebSize         = 0;
    idle->ExitEvent       = (EVENT*)0;
    idle->NameLock        = (SPINLOCK)SPINLOCK_INIT;
    memset(idle->Name, 0, sizeof(idle->Name));
    memcpy(idle->Name, "idle", sizeof("idle"));
    KiInitializeApcState(idle);
    idle->TrapFrame = (KTRAP_FRAME*)0;
    LIST_INIT(&idle->ProcessLink);
    LIST_INIT(&idle->SchedulerLink);

    void* kstack = MmAllocatePoolTag(KSTACK_SIZE, MM_POOL_TAG('k', 's', 't', 'k'));
    if (!kstack) {
        ObDereferenceObject(&idle->Header);
        return NTLSTATUS_NO_MEMORY;
    }
    idle->KernelStackBase = kstack;
    idle->KernelStackTop  = ((uintptr_t)kstack + KSTACK_SIZE) & ~(uintptr_t)15;
    idle->TssRsp0         = idle->KernelStackTop;

    uintptr_t* ksp = (uintptr_t*)idle->KernelStackTop;

    // Kernel stacks may come from allocators with weaker alignment than the
    // MS x64 call ABI requires. Bias the top down before building synthetic
    // entry frames so optimized C code can use aligned stack spills safely.
    // KiIdleLoop is a C function reached by a synthetic RET, not by CALL.
    // Leave the fake return slot and the MS x64 32-byte home area below the
    // top of stack so optimized prologues may spill above RSP without writing
    // into the pool canary or footer.
    ksp -= 5;
    __stosq((unsigned __int64*)ksp, 0, 5);

    ksp -= sizeof(KCONTEXT) / sizeof(uintptr_t);
    KCONTEXT* ctx = (KCONTEXT*)ksp;
    __stosq((unsigned __int64*)ctx, 0, sizeof(KCONTEXT) / 8);
    ctx->Rip = (uint64_t)KiIdleLoop;
    idle->SavedContext = ctx;
    KeInitThreadFpuContext(idle);

    // The hand-built CPU idle thread is not created through KeCreateThread(),
    // so register it manually with the dedicated idle process while kernel worker
    // threads remain grouped under the system process.
    ObReferenceObject(&gIdleProcess->Header);
    uintptr_t idleFlags = KeAcquireSpinlock(&gIdleProcess->ThreadLock);
    KeListInsertTailChecked(&gIdleProcess->ThreadList, &idle->ProcessLink);
    gIdleProcess->ThreadCount++;
    KeReleaseSpinlock(&gIdleProcess->ThreadLock, idleFlags);

    cpu->IdleThread = idle;
    cpu->CurrentThread = idle;
    cpu->OutgoingThread = (THREAD*)0;
    idle->SwitchLockCookie = 0x202ULL;
    return NTLSTATUS_SUCCESS;
}

// Initializes scheduler.
void KeInitializeScheduler(void) {
    KiInitializeSchedulerGlobal();

    KCPU* cpu = KeGetCurrentCpu();
    KSTATUS s = KiCreateIdleThreadForCpu(cpu);
    if (!KSUCCESS(s))
        KePanic(KE_PANIC_CODE_INIT, "scheduler: failed to create CPU%u idle thread, status=%u\n",
                (unsigned)cpu->CpuIndex, s);

    NtlAtomicExchange32(&cpu->State, KeCpuOnline);
}

// Safely release the lock and clean up any dead threads natively post-switch.
void KeSchedulerPostSwitch(void) {
    KCPU* cpu = KeGetCurrentCpu();
    THREAD* cur = cpu->CurrentThread;
    KeReleaseSpinlock(&gSchedLock, KiPostSwitchUnlockCookie(cur));

    if (cpu->OutgoingThread && cpu->OutgoingThread->State == ThreadDead) {
        THREAD* dead = cpu->OutgoingThread;
        cpu->OutgoingThread = (THREAD*)0;
        dead->RunningCpu = KE_INVALID_CPU;

        // Forced process termination can bypass KeExitThread. Finalize thread
        // ownership after the dead thread switches out, before releasing the
        // scheduler lifetime reference.
        KeFinalizeThreadExit(dead);
        KiMutexRundownThread(dead);
        ObDereferenceObject(&dead->Header);
    }

    KiDeliverKernelApcs();
}

// Updates thread Affinity Mask.
void KeSetThreadAffinityMask(THREAD* thread, uint64_t affinityMask) {
    if (!thread || thread->IsIdleThread) return;

    if (!affinityMask)
        affinityMask = 1ULL;

    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    thread->AffinityMask = affinityMask;
    if (thread->Queued)
        KiPokeRunnableCpusLocked(thread);
    KeReleaseSpinlock(&gSchedLock, flags);
}

// Returns default Thread Affinity Mask.
uint64_t KeGetDefaultThreadAffinityMask(void) {
    return KiDefaultThreadAffinityMask();
}

// Implements the ready Thread operation.
void KeReadyThread(THREAD* thread) {
    if (!thread) return;

    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    KiInsertReadyLocked(thread);
    KiPokeRunnableCpusLocked(thread);

    KeReleaseSpinlock(&gSchedLock, flags);
}

// Decrements one thread's suspend count and queues it when the count reaches zero.
__declspec(dllexport) KSTATUS KeResumeThread(THREAD* thread,
                                             uint32_t* previousSuspendCountOut) {
    if (!thread || thread->IsIdleThread)
        return NTLSTATUS_INVALID_PARAMETER;

    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    if (thread->State == ThreadDead) {
        KeReleaseSpinlock(&gSchedLock, flags);
        return NTLSTATUS_OBJECT_DEAD;
    }

    uint32_t previous = thread->SuspendCount;
    if (previous != 0) {
        if (thread->State != ThreadSuspended || thread->Queued ||
            thread->RunningCpu != KE_INVALID_CPU) {
            KiSchedulerBug("KeResumeThread", thread,
                           "suspend count did not match scheduler state");
        }

        thread->SuspendCount = previous - 1;
        if (thread->SuspendCount == 0) {
            KiInsertReadyLocked(thread);
            KiPokeRunnableCpusLocked(thread);
        }
    }
    KeReleaseSpinlock(&gSchedLock, flags);

    if (previousSuspendCountOut)
        *previousSuspendCountOut = previous;
    return NTLSTATUS_SUCCESS;
}

// Cancels a thread that has never left its initial suspended state.
KSTATUS KeCancelSuspendedThread(THREAD* thread) {
    if (!thread || thread->IsIdleThread)
        return NTLSTATUS_INVALID_PARAMETER;

    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    if (thread->State != ThreadSuspended || thread->SuspendCount == 0 ||
        thread->Queued || thread->RunningCpu != KE_INVALID_CPU) {
        KeReleaseSpinlock(&gSchedLock, flags);
        return NTLSTATUS_INVALID_PARAMETER;
    }

    thread->SuspendCount = 0;
    thread->State = ThreadDead;
    KeReleaseSpinlock(&gSchedLock, flags);

    KeFinalizeThreadExit(thread);

    // Drop the scheduler-owned lifetime reference. The creator still owns the
    // reference returned by KeCreateThread().
    ObDereferenceObject(&thread->Header);
    return NTLSTATUS_SUCCESS;
}

// Implements the unready Thread operation.
void KeUnreadyThread(THREAD* thread) {
    if (!thread) return;
    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    if (thread->Queued) {
        KeListRemoveChecked(&thread->SchedulerLink);
        thread->Queued = 0;
    }
    KeReleaseSpinlock(&gSchedLock, flags);
}

// Implements the any Runnable For Cpu Locked operation.
static bool KiAnyRunnableForCpuLocked(KCPU* cpu) {
    if (!cpu) cpu = KeGetCurrentCpu();

    for (uint32_t p = 0; p < KE_REALTIME_PRIORITY_COUNT; p++) {
        LIST_ENTRY* head = &gRealtimeRunQueue[p];
        for (LIST_ENTRY* node = head->Fwd; node != head; node = node->Fwd) {
            THREAD* thread = CONTAINING_RECORD(node, THREAD, SchedulerLink);
            if (!KiThreadCanRunOnCpuLocked(thread, cpu))
                continue;
            if (KiShouldDeferToIdlePrimaryLocked(cpu, thread))
                continue;
            return true;
        }
    }

    for (int p = 0; p < PriorityCount; p++) {
        LIST_ENTRY* head = &gNormalRunQueue[p];
        for (LIST_ENTRY* node = head->Fwd; node != head; node = node->Fwd) {
            THREAD* thread = CONTAINING_RECORD(node, THREAD, SchedulerLink);
            if (!KiThreadCanRunOnCpuLocked(thread, cpu))
                continue;
            if (KiShouldDeferToIdlePrimaryLocked(cpu, thread))
                continue;
            return true;
        }
    }
    return false;
}

// Removes the first eligible thread from one ready queue.
static THREAD* KiPickFromQueueLocked(KCPU* cpu, LIST_ENTRY* head) {
    for (LIST_ENTRY* node = head->Fwd; node != head; node = node->Fwd) {
        THREAD* thread = CONTAINING_RECORD(node, THREAD, SchedulerLink);
        if (!KiThreadCanRunOnCpuLocked(thread, cpu))
            continue;
        if (KiShouldDeferToIdlePrimaryLocked(cpu, thread)) {
            gSchedStats.SmtSiblingDeferrals++;
            continue;
        }

        KeListRemoveChecked(node);
        thread->Queued = 0;
        KiAssertPickedThreadLocked("KiPickFromQueueLocked", thread);
        return thread;
    }
    return (THREAD*)0;
}

// Implements the pick Next For Cpu operation.
static THREAD* KiPickNextForCpu(KCPU* cpu) {
    if (!cpu) cpu = KeGetCurrentCpu();

    for (int p = (int)KE_REALTIME_PRIORITY_MAX; p >= 0; p--) {
        THREAD* thread = KiPickFromQueueLocked(cpu, &gRealtimeRunQueue[p]);
        if (thread) {
            gSchedStats.RealtimeDispatches++;
            return thread;
        }
    }

    for (int p = PriorityCount - 1; p >= 0; p--) {
        THREAD* thread = KiPickFromQueueLocked(cpu, &gNormalRunQueue[p]);
        if (thread)
            return thread;
    }

    return cpu->IdleThread;
}

// Implements the request Reschedule operation.
void KiRequestReschedule(void) {
    KCPU* cpu = KeGetCurrentCpu();
    if (cpu)
        NtlAtomicExchange32(&cpu->ReschedulePending, 1);
}

// Checks whether reschedule Pending.
bool KiIsReschedulePending(void) {
    KCPU* cpu = KeGetCurrentCpu();
    return cpu && NtlAtomicLoad32(&cpu->ReschedulePending) != 0;
}

// Implements the schedule operation.
static void KiSchedule(void) {
    KCPU* cpu = KeGetCurrentCpu();
    THREAD* cur = cpu->CurrentThread;
    THREAD* idle = cpu->IdleThread;
    if (!cur) return;

    NtlAtomicExchange32(&cpu->ReschedulePending, 0);

    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    KiArmPostSwitchUnlock(cur, flags);

    // The scheduler never converts process termination directly into
    // ThreadDead. Terminal state owns rundown/lifetime publication and is
    // entered only through KeExitThread at a safe execution boundary.
    // If the current thread is still runnable and nobody else is ready, keep
    // executing it rather than alternating with idle on every timer tick.
    if (cur != idle && cur->State == ThreadRunning && !KiAnyRunnableForCpuLocked(cpu)) {
        if (cur->QuantumRemaining == 0)
            KiRefillQuantumLocked(cur);
        gSchedStats.KeepCurrentDecisions++;
        KeReleaseSpinlock(&gSchedLock, flags);
        return;
    }

    THREAD* next = KiPickNextForCpu(cpu);

    // A voluntary reschedule must not let lower-ranked work displace a running
    // real-time thread. Equal-ranked peers may still rotate at the quantum
    // boundary, and higher-ranked real-time work may preempt immediately.
    if (cur != idle && cur->State == ThreadRunning &&
        cur->EffectiveSchedulingClass == KeSchedulingRealtime &&
        next != cur &&
        !KiSchedulerThreadOutranks(next, cur) &&
        !KiRanksEqual(next->EffectiveSchedulingClass, next->EffectivePriority,
                      cur->EffectiveSchedulingClass, cur->EffectivePriority)) {
        if (next != idle)
            KiInsertReadyLocked(next);
        if (cur->QuantumRemaining == 0)
            KiRefillQuantumLocked(cur);
        gSchedStats.KeepCurrentDecisions++;
        KeReleaseSpinlock(&gSchedLock, flags);
        return;
    }

    if (next == idle) gSchedStats.IdleSelections++;
    if (next == cur) {
        if (cur != idle)
            KiSchedulerBug("KiSchedule", cur, "current thread unexpectedly appeared in ready queue");
        KeReleaseSpinlock(&gSchedLock, flags);
        return;
    }

    KiAssertPickedThreadLocked("KiSchedule", next);
    next->State    = ThreadRunning;
    KiRefillQuantumLocked(next);
    cpu->CurrentThread = next;
    next->RunningCpu = cpu->CpuIndex;
    next->LastCpu = cpu->CpuIndex;

    // Only re-queue the outgoing thread if it is still in the running state.
    // If KeBlockThread set State = ThreadWaiting, SchedulerLink was already
    // inserted into the sync object's wait list - inserting it here too would
    // corrupt both lists simultaneously (same intrusive node in two lists).
    // If KeExitThread set State = ThreadDead, the thread must not be re-queued.
    if (cur != idle && cur->State == ThreadRunning) {
        cur->RunningCpu = KE_INVALID_CPU;
        cur->LastCpu = cpu->CpuIndex;
        cur->State = ThreadReady;
        KiInsertReadyLocked(cur);
    } else if (cur) {
        cur->RunningCpu = KE_INVALID_CPU;
        cur->LastCpu = cpu->CpuIndex;
    }
    cpu->OutgoingThread = cur;
    cpu->ContextSwitches++;
    gSchedStats.ContextSwitches++;
    KeSwitchContext(cur, next);
    KeSchedulerPostSwitch();
}

// Implements the schedule From Interrupt operation.
void KiScheduleFromInterrupt(void) {
    if (!KiIsReschedulePending()) return;
    KiSchedule();
}

// Implements the scheduler Mark Thread Waiting operation.
void KiSchedulerMarkThreadWaiting(THREAD* thread) {
    if (!thread) return;

    KCPU* cpu = KeGetCurrentCpu();
    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);

    if (thread->IsIdleThread)
        KiSchedulerBug("KiSchedulerMarkThreadWaiting", thread, "idle thread cannot block");
    if (thread->Queued)
        KiSchedulerBug("KiSchedulerMarkThreadWaiting", thread, "queued thread cannot block");
    if (thread != cpu->CurrentThread)
        KiSchedulerBug("KiSchedulerMarkThreadWaiting", thread, "only the current thread can prepare a blocking wait");
    if (thread->RunningCpu != cpu->CpuIndex)
        KiSchedulerBug("KiSchedulerMarkThreadWaiting", thread, "thread running CPU does not match current CPU");
    if (thread->State != ThreadRunning)
        KiSchedulerBug("KiSchedulerMarkThreadWaiting", thread, "thread was not running before wait");
    if (thread->TimeoutArmed ||
        thread->TimeoutLink.Fwd != &thread->TimeoutLink ||
        thread->TimeoutLink.Bwd != &thread->TimeoutLink) {
        KiSchedulerBug("KiSchedulerMarkThreadWaiting", thread, "thread already has a timeout");
    }

    thread->State = ThreadWaiting;
    KeReleaseSpinlock(&gSchedLock, flags);
}

// Cancels one armed deadline while the scheduler lock is held.
static void KiCancelThreadTimeoutLocked(THREAD* thread) {
    if (!thread) return;

    if (thread->TimeoutArmed) {
        if (thread->TimeoutLink.Fwd == &thread->TimeoutLink ||
            thread->TimeoutLink.Bwd == &thread->TimeoutLink) {
            KiSchedulerBug("KiCancelThreadTimeoutLocked", thread, "armed timeout was not linked");
        }
        KeListRemoveChecked(&thread->TimeoutLink);
        thread->TimeoutArmed = 0;
    } else if (thread->TimeoutLink.Fwd != &thread->TimeoutLink ||
               thread->TimeoutLink.Bwd != &thread->TimeoutLink) {
        KiSchedulerBug("KiCancelThreadTimeoutLocked", thread, "unarmed timeout remained linked");
    }

    thread->TimeoutTick = 0;
    thread->TimeoutStatus = NTLSTATUS_SUCCESS;
}

// Implements the scheduler Note Voluntary Block operation.
void KiSchedulerNoteVoluntaryBlock(THREAD* thread) {
    if (!thread) return;
    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    gSchedStats.VoluntaryBlocks++;
    thread->QuantumRemaining = 0;
    KeReleaseSpinlock(&gSchedLock, flags);
}

// Completes a wait while the scheduler lock is held.
static void KiWakeThreadLocked(THREAD* thread, KSTATUS status) {
    if (!thread) return;

    // Complete the wait while holding the scheduler lock. On SMP a wake can
    // race the prepare->commit window on another CPU: the target thread is
    // already on the object's wait list and has State=ThreadWaiting, but it is
    // still executing with RunningCpu set until it actually calls KeYield().
    // That case must complete in-place, not enqueue the still-running thread.
    // Once another path has completed the wait, a stale abort or duplicate
    // signal must not replace the result that made the thread runnable.
    if (thread->State != ThreadWaiting)
        return;

    KiCancelThreadTimeoutLocked(thread);

    thread->WaitObject    = (OBJECT*)0;
    thread->WaitStatus    = status;
    thread->WaitListHead  = (LIST_ENTRY*)0;
    thread->AlertableWait = 0;

    thread->CpuBurstTicks = 0;
    thread->QuantaConsumed = 0;
    thread->QuantumRemaining = 0;

    if (thread->RunningCpu != KE_INVALID_CPU) {
        // Prepare->commit race: the waiter is still running somewhere. Its
        // KeCommitBlockThread() will see State != ThreadWaiting and return the
        // completed wait status without ever entering the ready queue.
        thread->State = ThreadRunning;
        return;
    }

    if (thread->BaseSchedulingClass == KeSchedulingNormal &&
        thread->BasePriority < PriorityHigh) {
        thread->BasePriority++;
    }

    KE_SCHEDULING_CLASS schedulingClass;
    uint32_t priority;
    KiComputeEffectiveRankLocked(thread, &schedulingClass, &priority);
    (void)KiApplyEffectiveRankLocked(thread, schedulingClass, priority);
    KiInsertReadyLocked(thread);
    KiPokeRunnableCpusLocked(thread);
}

// Wakes a thread after first removing its scheduler-link wait membership.
void KeUnblockThreadFromWaitListWithStatus(THREAD* thread, KSTATUS status) {
    if (!thread) return;
    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    if (thread->State == ThreadWaiting && thread->WaitListHead) {
        KeListRemoveChecked(&thread->SchedulerLink);
        thread->WaitListHead = (LIST_ENTRY*)0;
    }
    KiWakeThreadLocked(thread, status);
    KeReleaseSpinlock(&gSchedLock, flags);
}

// Implements the scheduler Wake Thread operation.
void KiSchedulerWakeThread(THREAD* thread, KSTATUS status) {
    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    KiWakeThreadLocked(thread, status);
    KeReleaseSpinlock(&gSchedLock, flags);
}

// Detaches one expired deadline while the scheduler lock is held.
static THREAD* KiTakeExpiredTimeoutLocked(uint64_t nowTick, KSTATUS* statusOut) {
    LIST_ENTRY* head = &gTimeoutList;
    for (LIST_ENTRY* node = head->Fwd; node != head; node = node->Fwd) {
        THREAD* thread = CONTAINING_RECORD(node, THREAD, TimeoutLink);
        if (!thread->TimeoutArmed)
            KiSchedulerBug("KiTakeExpiredTimeoutLocked", thread, "timeout list contains an unarmed thread");
        if (thread->State != ThreadWaiting)
            KiSchedulerBug("KiTakeExpiredTimeoutLocked", thread, "timeout list contains a non-waiting thread");
        if (thread->TimeoutTick > nowTick)
            continue;

        KSTATUS timeoutStatus = thread->TimeoutStatus;
        KeListRemoveChecked(&thread->TimeoutLink);
        thread->TimeoutArmed = 0;
        thread->TimeoutTick = 0;
        thread->TimeoutStatus = NTLSTATUS_SUCCESS;
        ObReferenceObject(&thread->Header);
        if (statusOut) *statusOut = timeoutStatus;
        return thread;
    }

    return (THREAD*)0;
}

// Reports whether one deadline is ready while the scheduler lock is held.
static bool KiHasExpiredTimeoutLocked(uint64_t nowTick) {
    LIST_ENTRY* head = &gTimeoutList;
    for (LIST_ENTRY* node = head->Fwd; node != head; node = node->Fwd) {
        THREAD* thread = CONTAINING_RECORD(node, THREAD, TimeoutLink);
        if (thread->TimeoutTick <= nowTick)
            return true;
    }
    return false;
}

// Resolves expired deadlines at passive level through each object's abort path.
static void KiTimeoutWorkRoutine(void* context) {
    (void)context;
    KeAssertIrql(PASSIVE_LEVEL);

    while (true) {
        KSTATUS timeoutStatus = NTLSTATUS_TIMEOUT;
        uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
        uint64_t nowTick = gSchedStats.Ticks;
        THREAD* thread = KiTakeExpiredTimeoutLocked(nowTick, &timeoutStatus);
        KeReleaseSpinlock(&gSchedLock, flags);
        if (!thread)
            return;

        (void)KeAbortThreadWait(thread, timeoutStatus);
        ObDereferenceObject(&thread->Header);
    }
}

// Checks whether this CPU owns global scheduler tick duties.
static bool KiCpuOwnsGlobalSchedulerTick(const KCPU* cpu) {
    // Periodic timer interrupts are per-CPU, but sleep deadlines, waitable timers,
    // input repeat, and priority aging are global scheduler services. Keep the
    // global clock single-sourced on the boot processor so enabling AP timers does not
    // make those services run once per online CPU.
    return cpu && cpu->IsBootProcessor;
}

// Ages ready normal-class threads after enough global ticks have elapsed.
static void KiAgeReadyThreadsLocked(void) {
    // Real-time queues are static and never participate in aging. Normal
    // threads promoted by aging get a fresh CPU-burst history because they
    // earned the boost by waiting rather than by burning CPU.
    gAgingCounter++;
    if (gAgingCounter < SCHED_AGING_INTERVAL_TICKS)
        return;

    gAgingCounter = 0;
    for (int p = 0; p < PriorityHigh; p++) {
        while (!LIST_IS_EMPTY(&gNormalRunQueue[p])) {
            LIST_ENTRY* node = KeListPopHeadChecked(&gNormalRunQueue[p]);
            THREAD* thread = CONTAINING_RECORD(node, THREAD, SchedulerLink);
            thread->Queued = 0;
            thread->BasePriority = PriorityHigh;
            thread->CpuBurstTicks = 0;
            thread->QuantaConsumed = 0;
            thread->QuantumRemaining = 0;
            KE_SCHEDULING_CLASS schedulingClass;
            uint32_t priority;
            KiComputeEffectiveRankLocked(thread, &schedulingClass, &priority);
            thread->EffectiveSchedulingClass = schedulingClass;
            thread->EffectivePriority = priority;
            gSchedStats.AgingPromotions++;
            KiInsertReadyLocked(thread);
        }
    }
}

// Runs scheduler work that must be single-sourced for the whole system.
static int KiRunGlobalSchedulerTickLocked(void) {
    gSchedStats.Ticks++;
    KiAgeReadyThreadsLocked();
    return KiHasExpiredTimeoutLocked(gSchedStats.Ticks);
}

// Services the current CPU's local scheduler timer interrupt.
static int KiRunLocalSchedulerTickLocked(KCPU* cpu, THREAD* cur, THREAD* idle) {
    int requestReschedule = 0;

    if (cur == idle) {
        if (KiAnyRunnableForCpuLocked(cpu))
            requestReschedule = 1;
    } else if (KeThreadExitInProgress(cur)) {
        // Terminal teardown is interruptible but not preemptible. Its explicit
        // final KeYield performs the permanent switch after completion state is
        // published.
        requestReschedule = 0;
    } else if (KiThreadIsKilledLocked(cur)) {
        requestReschedule = 1;
    } else {
        if (cur->EffectiveSchedulingClass == KeSchedulingNormal) {
            cur->CpuBurstTicks++;
            cur->QuantaConsumed = cur->CpuBurstTicks;

            if (cur->CpuBurstTicks >= SCHED_DEMOTE_AFTER_BURST_TICKS) {
                cur->CpuBurstTicks = 0;
                cur->QuantaConsumed = 0;
                if (cur->BasePriority > PriorityLow) {
                    cur->BasePriority--;
                    KE_SCHEDULING_CLASS schedulingClass;
                    uint32_t priority;
                    KiComputeEffectiveRankLocked(cur, &schedulingClass, &priority);
                    (void)KiApplyEffectiveRankLocked(cur, schedulingClass, priority);
                    gSchedStats.Demotions++;
                    gSchedStats.CpuBurstDemotions++;
                }
            }
        }

        if (cur->QuantumRemaining == 0)
            KiRefillQuantumLocked(cur);

        if (cur->QuantumRemaining > 0)
            cur->QuantumRemaining--;

        if (cur->QuantumRemaining == 0) {
            int havePeer = cur->EffectiveSchedulingClass == KeSchedulingRealtime
                ? KiAnyRealtimePeerForCpuLocked(cpu, cur->EffectivePriority)
                : KiAnyRunnableForCpuLocked(cpu);
            if (havePeer) {
                gSchedStats.QuantumExpirations++;
                requestReschedule = 1;
            } else {
                // No eligible peer is ready, so keep the CPU hot and refill the
                // budget instead of turning a periodic tick into idle chatter.
                KiRefillQuantumLocked(cur);
                gSchedStats.KeepCurrentDecisions++;
            }
        }
    }

    return requestReschedule;
}

// Implements the scheduler Tick operation.
void KeSchedulerTick(void) {
    KCPU* cpu = KeGetCurrentCpu();
    THREAD* cur = cpu->CurrentThread;
    THREAD* idle = cpu->IdleThread;
    if (!cur) return;

    int ownsGlobalTick = KiCpuOwnsGlobalSchedulerTick(cpu);
    int queueTimeoutWork = 0;

    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    KiArmPostSwitchUnlock(cur, flags);

    if (ownsGlobalTick)
        queueTimeoutWork = KiRunGlobalSchedulerTickLocked();

    int requestReschedule = KiRunLocalSchedulerTickLocked(cpu, cur, idle);

    KeReleaseSpinlock(&gSchedLock, flags);

    if (ownsGlobalTick) {
        if (queueTimeoutWork)
            (void)KeQueueWorkItem(&gTimeoutWorkItem);
        KeTimerTick();
        KeDisplayTick();
    }

    if (requestReschedule)
        KiRequestReschedule();
}

// Converts one requested delay to scheduler ticks while reserving one phase guard tick.
static uint64_t KiTimeoutTicksForMilliseconds(uint64_t milliseconds) {
    uint64_t ticks = milliseconds / SCHED_TICK_MS;
    if ((milliseconds % SCHED_TICK_MS) != 0)
        ticks++;

    // The scheduler tick count identifies completed quanta, not the current
    // position inside one quantum. One additional tick guarantees that a
    // timeout armed immediately before an interrupt cannot expire early.
    if (ticks < UINT64_MAX)
        ticks++;
    return ticks ? ticks : 1;
}

// Returns the saturated scheduler tick at which one finite timeout may expire.
static uint64_t KiTimeoutWakeTickLocked(uint64_t milliseconds) {
    uint64_t ticks = KiTimeoutTicksForMilliseconds(milliseconds);
    uint64_t now = gSchedStats.Ticks;
    if (ticks == UINT64_MAX || ticks > UINT64_MAX - now)
        return UINT64_MAX;
    return now + ticks;
}

// Arms a timeout for the current thread after a wait primitive has prepared it to block.
KSTATUS KeArmCurrentThreadTimeout(uint64_t milliseconds, KSTATUS timeoutStatus) {
    if (milliseconds == 0)
        return NTLSTATUS_INVALID_PARAMETER;
    if (KeGetCurrentIrql() > DISPATCH_LEVEL)
        return NTLSTATUS_INVALID_IRQL;

    THREAD* self = KeGetCurrentThread();
    if (!self || self->IsIdleThread)
        return NTLSTATUS_INVALID_PARAMETER;

    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    if (self->State != ThreadWaiting ||
        self->RunningCpu != KeGetCurrentCpu()->CpuIndex ||
        self->TimeoutArmed ||
        self->TimeoutLink.Fwd != &self->TimeoutLink ||
        self->TimeoutLink.Bwd != &self->TimeoutLink) {
        KeReleaseSpinlock(&gSchedLock, flags);
        return NTLSTATUS_INVALID_PARAMETER;
    }

    self->TimeoutTick = KiTimeoutWakeTickLocked(milliseconds);
    self->TimeoutStatus = timeoutStatus;
    self->TimeoutArmed = 1;
    KeListInsertTailChecked(&gTimeoutList, &self->TimeoutLink);
    KeReleaseSpinlock(&gSchedLock, flags);
    return NTLSTATUS_SUCCESS;
}

// Implements the delay Thread operation.
__declspec(dllexport) KSTATUS KeDelayThread(uint64_t milliseconds) {
    if (milliseconds == 0) {
        KeYield();
        return NTLSTATUS_SUCCESS;
    }

    if (!KeCanWait())
        return NTLSTATUS_INVALID_IRQL;

    THREAD* self = KeGetCurrentThread();
    if (!self || self->IsIdleThread)
        return NTLSTATUS_INVALID_PARAMETER;

    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);

    if (self->Queued || self->State != ThreadRunning ||
        self->RunningCpu != KeGetCurrentCpu()->CpuIndex ||
        self->TimeoutArmed ||
        self->TimeoutLink.Fwd != &self->TimeoutLink ||
        self->TimeoutLink.Bwd != &self->TimeoutLink) {
        KeReleaseSpinlock(&gSchedLock, flags);
        return NTLSTATUS_INVALID_PARAMETER;
    }

    uint64_t wakeTick = KiTimeoutWakeTickLocked(milliseconds);

    self->State = ThreadWaiting;
    self->WaitObject = (OBJECT*)0;
    self->WaitStatus = NTLSTATUS_PENDING;
    self->WaitListHead = (LIST_ENTRY*)0;
    self->AlertableWait = 0;
    self->TimeoutTick = wakeTick;
    self->TimeoutStatus = NTLSTATUS_SUCCESS;
    self->TimeoutArmed = 1;
    KeListInsertTailChecked(&gTimeoutList, &self->TimeoutLink);

    KeReleaseSpinlock(&gSchedLock, flags);
    return KeCommitBlockThread(self, 0);
}

// Implements the yield operation.
__declspec(dllexport) void KeYield(void) {
    KeAssertIrql(PASSIVE_LEVEL);
    KiSchedule();
}