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

#include <stddef.h>
#include <stdbool.h>
#include <intrin.h>
#pragma intrinsic(_enable)
#include <kernel/ke_sched.h>
#include <kernel/ke_cpu.h>
#include <kernel/ke_process.h>
#include <kernel/ke_thread.h>
#include <kernel/ke_spinlock.h>
#include <kernel/ke_irql.h>
#include <kernel/ke_apc.h>
#include <kernel/ke_debug.h>
#include <kernel/ke_list.h>
#include <kernel/mm_virt.h>
#include <kernel/hal/hal_apic.h>

static LIST_ENTRY gRunQueue[PriorityCount];
static LIST_ENTRY gSleepList;
static SPINLOCK   gSchedLock;

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

// Implements the default Thread Affinity Mask operation.
static uint64_t KiDefaultThreadAffinityMask(void) {
    // Normal kernel, user, and driver-created threads are BSP-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 int KiThreadIsKilledLocked(const THREAD* thread) {
    return thread && !thread->IsIdleThread && thread->Owner && KeIsProcessTerminating(thread->Owner);
}

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

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

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

    for (uint32_t i = 0; i < gCpuCount; i++) {
        KCPU* cpu = &gCpus[i];
        if (!cpu->Present || !cpu->Started || !cpu->Online)
            continue;
        if (!cpu->IsPrimaryThread)
            continue;
        if (!KiCpuIsIdleLocked(cpu))
            continue;
        if (KiThreadCanRunOnCpuLocked(thread, cpu))
            return 1;
    }
    return 0;
}

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

    // 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->Online)
        return;

    _InterlockedExchange(&target->ReschedulePending, 1);

    KCPU* self = KeGetCurrentCpu();
    if (self && target->CpuIndex != self->CpuIndex) {
        HalApicSendRescheduleIpi(target->LapicId);
    }
}

// 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->Present || !cpu->Started || !cpu->Online)
            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 && thread->Priority > cur->Priority) {
            gSchedStats.PriorityPreemptions++;
            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->Present || !cpu->Started || !cpu->Online)
                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(THREAD_PRIORITY priority) {
    switch (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;

    uint32_t quantum = KiBaseQuantumForPriority(thread->Priority);

    // 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->Priority == 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 ready Locked.
static void KiInsertReadyLocked(THREAD* thread) {
    KiAssertReadyQueueCandidateLocked("KiInsertReadyLocked", thread);

    int p = (int)thread->Priority;
    if (p < 0) p = 0;
    if (p >= PriorityCount) p = PriorityCount - 1;

    thread->State = ThreadReady;
    thread->Queued = 1;
    ListInsertTail(&gRunQueue[p], &thread->SchedulerLink);
}

static int 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) {
    for (;;) {
        KCPU* cpu = KeGetCurrentCpu();
        cpu->SchedulerTicks++;

        // AP idle threads sleep until an interrupt or reschedule IPI wakes the
        // CPU. Ready paths set ReschedulePending and target the CPU's LAPIC;
        // 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.
        _enable();
        __halt();
    }
}

// Initializes scheduler Global.
void KiInitializeSchedulerGlobal(void) {
    for (int i = 0; i < PriorityCount; i++) LIST_INIT(&gRunQueue[i]);
    LIST_INIT(&gSleepList);
    gSchedLock = (SPINLOCK)SPINLOCK_INIT;

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

    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 STATUS_INVALID_PARAMETER;
    if (cpu->IdleThread)
        return STATUS_SUCCESS;

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

    idle->Header.Destroy = (void *)0;
    idle->ThreadId        = cpu->CpuIndex ? (0x80000000u | cpu->CpuIndex) : 0;
    idle->State           = ThreadRunning;
    idle->Priority        = PriorityLow;
    idle->QuantaConsumed  = 0;
    idle->QuantumTicks     = 1;
    idle->QuantumRemaining = 1;
    idle->CpuBurstTicks    = 0;
    idle->LastRunTicks     = 0;
    idle->RunningCpu       = cpu->CpuIndex;
    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      = STATUS_SUCCESS;
    idle->WaitListHead    = (LIST_ENTRY*)0;
    idle->AlertableWait   = 0;
    idle->SleepUntilTick  = 0;
    idle->TimeoutStatus   = STATUS_SUCCESS;
    idle->UserStackBase   = (void*)0;
    idle->Teb             = (void*)0;
    idle->TebSize         = 0;
    idle->ExitEvent       = (EVENT*)0;
    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 STATUS_NO_MEMORY;
    }
    idle->KernelStackBase = kstack;
    idle->KernelStackTop  = (uintptr_t)kstack + KSTACK_SIZE;
    idle->TssRsp0         = idle->KernelStackTop;

    uintptr_t* ksp = (uintptr_t *)idle->KernelStackTop;
    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 idle process. This keeps PID 0 visible
    // as the real idle process while kernel worker threads live under PID 1.
    ObReferenceObject(&gIdleProcess->Header);
    uintptr_t idleFlags = KeAcquireSpinlock(&gIdleProcess->HandleLock);
    ListInsertTail(&gIdleProcess->ThreadList, &idle->ProcessLink);
    gIdleProcess->ThreadCount++;
    KeReleaseSpinlock(&gIdleProcess->HandleLock, idleFlags);

    cpu->IdleThread = idle;
    cpu->CurrentThread = idle;
    cpu->OutgoingThread = (THREAD*)0;
    idle->SwitchLockCookie = 0x202ULL;
    return STATUS_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);
}

// 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));

    KiDeliverKernelApcs();

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

// 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);
}

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


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

    for (int p = 0; p < PriorityCount; p++) {
        LIST_ENTRY* head = &gRunQueue[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 1;
        }
    }
    return 0;
}

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

    for (int p = PriorityCount - 1; p >= 0; p--) {
        LIST_ENTRY* head = &gRunQueue[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)) {
                gSchedStats.SmtSiblingDeferrals++;
                continue;
            }

            ListRemove(node);
            thread->Queued = 0;
            KiAssertPickedThreadLocked("KiPickNextForCpu", thread);
            return thread;
        }
    }

    return cpu->IdleThread;
}


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

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

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

    _InterlockedExchange(&cpu->ReschedulePending, 0);

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

    // The currently running thread is not present in a run queue. If it is
    // still runnable and nobody else is ready, keep executing it. Otherwise a
    // lone CPU-bound thread gets alternated with the idle thread on every timer
    // tick, effectively turning a periodic preemption check into a 50% throttle.
    if (cur != idle && KiThreadIsKilledLocked(cur)) {
        cur->State = ThreadDead;
    }

    if (cur != idle && cur->State == ThreadRunning && !KiAnyRunnableForCpuLocked(cpu)) {
        if (cur->QuantumRemaining == 0)
            KiRefillQuantumLocked(cur);
        gSchedStats.KeepCurrentDecisions++;
        KeReleaseSpinlock(&gSchedLock, flags);
        return;
    }

    THREAD* next = KiPickNextForCpu(cpu);
    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");

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

// 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.
    if (thread->State != ThreadWaiting) {
        thread->WaitStatus = status;
        return;
    }

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

    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->Priority < PriorityHigh) {
        thread->Priority = (THREAD_PRIORITY)(thread->Priority + 1);
    }

    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) {
        ListRemove(&thread->SchedulerLink);
        thread->WaitListHead = (LIST_ENTRY*)0;
        thread->SleepUntilTick = 0;
        thread->TimeoutStatus = STATUS_SUCCESS;
    }
    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);
}

// Wakes threads whose scheduler-tick sleep deadline has expired.
static void KiWakeExpiredSleepersLocked(uint64_t nowTick) {
    LIST_ENTRY* node = gSleepList.Fwd;
    while (node != &gSleepList) {
        LIST_ENTRY* next = node->Fwd;
        THREAD* thread = CONTAINING_RECORD(node, THREAD, SchedulerLink);

        if (thread->State == ThreadWaiting &&
            thread->WaitListHead == &gSleepList &&
            thread->SleepUntilTick <= nowTick) {
            ListRemove(node);
            thread->WaitListHead = (LIST_ENTRY*)0;
            KiWakeThreadLocked(thread, thread->TimeoutStatus);
        }

        node = next;
    }
}

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

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

    gSchedStats.Ticks++;
    KiWakeExpiredSleepersLocked(gSchedStats.Ticks);

    int requestReschedule = 0;

    if (cur == idle) {
        if (KiAnyRunnableForCpuLocked(cpu))
            requestReschedule = 1;
    } else if (KiThreadIsKilledLocked(cur)) {
        requestReschedule = 1;
    } else {
        cur->CpuBurstTicks++;
        cur->QuantaConsumed = cur->CpuBurstTicks;

        if (cur->CpuBurstTicks >= SCHED_DEMOTE_AFTER_BURST_TICKS) {
            cur->CpuBurstTicks = 0;
            cur->QuantaConsumed = 0;
            if (cur->Priority > PriorityLow) {
                cur->Priority = (THREAD_PRIORITY)(cur->Priority - 1);
                gSchedStats.Demotions++;
                gSchedStats.CpuBurstDemotions++;
            }
        }

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

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

        if (cur->QuantumRemaining == 0) {
            if (KiAnyRunnableForCpuLocked(cpu)) {
                gSchedStats.QuantumExpirations++;
                requestReschedule = 1;
            } else {
                // No peer is ready, so keep the CPU hot and refill the budget.
                // This avoids reintroducing the old every-other-tick idle dance
                // as a pointless reschedule-pending chatter loop.
                KiRefillQuantumLocked(cur);
                gSchedStats.KeepCurrentDecisions++;
            }
        }
    }

    // Periodic aging pass to prevent indefinite starvation of PriorityLow and
    // PriorityNormal threads. Ready threads promoted by aging get a fresh CPU
    // burst history; they earned the boost by waiting, not by burning CPU.
    gAgingCounter++;
    if (gAgingCounter >= SCHED_AGING_INTERVAL_TICKS) {
        gAgingCounter = 0;
        for (int p = 0; p < PriorityHigh; p++) {
            while (!LIST_IS_EMPTY(&gRunQueue[p])) {
                LIST_ENTRY* node = ListPopHead(&gRunQueue[p]);
                THREAD* t = CONTAINING_RECORD(node, THREAD, SchedulerLink);
                t->Priority = PriorityHigh;
                t->CpuBurstTicks = 0;
                t->QuantaConsumed = 0;
                t->QuantumRemaining = 0;
                gSchedStats.AgingPromotions++;
                ListInsertTail(&gRunQueue[PriorityHigh], &t->SchedulerLink);
            }
        }
    }

    KeReleaseSpinlock(&gSchedLock, flags);

    if (requestReschedule)
        KiRequestReschedule();
}


// 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 STATUS_INVALID_PARAMETER;
    if (!KeCanWait())
        return STATUS_INVALID_IRQL;

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

    uint64_t ticks = milliseconds / SCHED_TICK_MS;
    if ((milliseconds % SCHED_TICK_MS) != 0)
        ticks++;
    if (ticks == 0)
        ticks = 1;

    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    if (self->State != ThreadWaiting || self->WaitListHead != (LIST_ENTRY*)0 ||
        self->RunningCpu != KeGetCurrentCpu()->CpuIndex) {
        KeReleaseSpinlock(&gSchedLock, flags);
        return STATUS_INVALID_PARAMETER;
    }

    uint64_t now = gSchedStats.Ticks;
    uint64_t wakeTick = now + ticks;
    if (wakeTick < now)
        wakeTick = UINT64_MAX;

    self->WaitListHead = &gSleepList;
    self->SleepUntilTick = wakeTick;
    self->TimeoutStatus = timeoutStatus;
    ListInsertTail(&gSleepList, &self->SchedulerLink);
    KeReleaseSpinlock(&gSchedLock, flags);
    return STATUS_SUCCESS;
}

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

    if (!KeCanWait())
        return STATUS_INVALID_IRQL;

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

    uint64_t ticks = milliseconds / SCHED_TICK_MS;
    if ((milliseconds % SCHED_TICK_MS) != 0)
        ticks++;
    if (ticks == 0)
        ticks = 1;

    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);

    if (self->Queued || self->State != ThreadRunning ||
        self->RunningCpu != KeGetCurrentCpu()->CpuIndex) {
        KeReleaseSpinlock(&gSchedLock, flags);
        return STATUS_INVALID_PARAMETER;
    }

    uint64_t now = gSchedStats.Ticks;
    uint64_t wakeTick = now + ticks;
    if (wakeTick < now)
        wakeTick = UINT64_MAX;

    self->State = ThreadWaiting;
    self->WaitObject = (OBJECT*)0;
    self->WaitStatus = STATUS_PENDING;
    self->WaitListHead = &gSleepList;
    self->AlertableWait = 0;
    self->SleepUntilTick = wakeTick;
    self->TimeoutStatus = STATUS_SUCCESS;
    ListInsertTail(&gSleepList, &self->SchedulerLink);

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

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