// ke_sched.c - MLFQ scheduler
#include <stddef.h>
#include <stdbool.h>
#include <ke_sched.h>
#include <ke_process.h>
#include <ke_thread.h>
#include <ke_spinlock.h>
#include <ke_irql.h>
#include <ke_apc.h>
#include <ke_list.h>
#include <mm_virt.h>

static LIST_ENTRY gRunQueue[PriorityCount];
static SPINLOCK   gSchedLock;

static THREAD* gCurrentThread;
static THREAD* gIdleThread;
static THREAD* gOutgoingThread;
static uint32_t gAgingCounter = 0;
static volatile long gReschedulePending = 0;
static KE_SCHED_STATS gSchedStats;

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

static uint32_t KiComputeQuantumLocked(THREAD* thread) {
    if (!thread || thread == gIdleThread)
        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;
}

static void KiRefillQuantumLocked(THREAD* thread) {
    if (!thread) return;
    thread->QuantumTicks = KiComputeQuantumLocked(thread);
    thread->QuantumRemaining = thread->QuantumTicks;
}

THREAD* KeGetCurrentThread(void)      { return gCurrentThread; }
void    KeSetCurrentThread(THREAD* t) { gCurrentThread = t; }

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

static void KiIdleLoop(void) {
    while (true) { __halt(); }
}

void KeInitializeScheduler(void) {
    for (int i = 0; i < PriorityCount; i++) LIST_INIT(&gRunQueue[i]);
    gSchedLock = (SPINLOCK)SPINLOCK_INIT;

    PROCESS* idleProc;
    KSTATUS s = KeCreateProcess("idle", &idleProc);
    if (!KSUCCESS(s)) return;

    PROCESS* systemProc;
    s = KeCreateSystemProcess(&systemProc);
    if (!KSUCCESS(s)) return;
    (void)systemProc;

    THREAD* idle = (THREAD *)ObAllocateObject(ObjectTypeThread, sizeof(THREAD));
    idle->Header.Destroy = (void *)0;
    idle->ThreadId        = 0;
    idle->State           = ThreadReady;
    idle->Priority        = PriorityLow;
    idle->QuantaConsumed  = 0;
    idle->QuantumTicks     = 1;
    idle->QuantumRemaining = 1;
    idle->CpuBurstTicks    = 0;
    idle->LastRunTicks     = 0;
    idle->Owner           = idleProc;
    idle->WaitObject      = (OBJECT *)0;
    idle->WaitStatus      = STATUS_SUCCESS;
    idle->WaitListHead    = (LIST_ENTRY*)0;
    idle->AlertableWait   = 0;
    idle->SchedLockFlags  = 0x202ULL; // IF=1
    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 = MmAllocatePool(KSTACK_SIZE);
    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(&idleProc->Header);
    uintptr_t idleFlags = KeAcquireSpinlock(&idleProc->HandleLock);
    ListInsertTail(&idleProc->ThreadList, &idle->ProcessLink);
    idleProc->ThreadCount++;
    KeReleaseSpinlock(&idleProc->HandleLock, idleFlags);

    gIdleThread    = idle;
    gCurrentThread = idle;
}

// Safely release the lock and clean up any dead threads natively post-switch
void KeSchedulerPostSwitch(void) {
    THREAD* cur = KeGetCurrentThread();
    KeReleaseSpinlock(&gSchedLock, cur->SchedLockFlags);

    KiDeliverKernelApcs();

    if (gOutgoingThread && gOutgoingThread->State == ThreadDead) {
        THREAD* dead = gOutgoingThread;
        gOutgoingThread = (THREAD*)0;
        ObDereferenceObject(&dead->Header);
    }
}

void KeReadyThread(THREAD* thread) {
    if (!thread) return;
    thread->State = ThreadReady;
    int p = (int)thread->Priority;
    if (p < 0) p = 0;
    if (p >= PriorityCount) p = PriorityCount - 1;

    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    ListInsertTail(&gRunQueue[p], &thread->SchedulerLink);

    THREAD* cur = gCurrentThread;
    if (!cur || cur == gIdleThread) {
        KiRequestReschedule();
    } else if (cur->State == ThreadRunning && thread->Priority > cur->Priority) {
        gSchedStats.PriorityPreemptions++;
        KiRequestReschedule();
    }

    KeReleaseSpinlock(&gSchedLock, flags);
}

void KeUnreadyThread(THREAD* thread) {
    if (!thread) return;
    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    ListRemove(&thread->SchedulerLink);
    KeReleaseSpinlock(&gSchedLock, flags);
}

static int KiAnyRunnableLocked(void) {
    for (int p = 0; p < PriorityCount; p++) {
        if (!LIST_IS_EMPTY(&gRunQueue[p]))
            return 1;
    }
    return 0;
}

static THREAD* KiPickNext(void) {
    for (int p = PriorityCount - 1; p >= 0; p--) {
        if (!LIST_IS_EMPTY(&gRunQueue[p])) {
            LIST_ENTRY* node = gRunQueue[p].Fwd;
            ListRemove(node);
            return CONTAINING_RECORD(node, THREAD, SchedulerLink);
        }
    }
    return gIdleThread;
}

void KiRequestReschedule(void) {
    gReschedulePending = 1;
}

int KiIsReschedulePending(void) {
    return gReschedulePending != 0;
}

static void KiSchedule(void) {
    THREAD* cur = gCurrentThread;
    if (!cur) return;

    gReschedulePending = 0;

    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    cur->SchedLockFlags = 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 != gIdleThread && cur->State == ThreadRunning && !KiAnyRunnableLocked()) {
        if (cur->QuantumRemaining == 0)
            KiRefillQuantumLocked(cur);
        gSchedStats.KeepCurrentDecisions++;
        KeReleaseSpinlock(&gSchedLock, flags);
        return;
    }

    THREAD* next = KiPickNext();
    if (next == gIdleThread) gSchedStats.IdleSelections++;
    if (next == cur) {
        // Only re-queue if still running normally; a ThreadWaiting/ThreadDead
        // caller has already placed SchedulerLink elsewhere (or is discarding
        // the thread).
        if (cur != gIdleThread && cur->State == ThreadRunning)
            ListInsertTail(&gRunQueue[cur->Priority], &cur->SchedulerLink);
        KeReleaseSpinlock(&gSchedLock, flags);
        return;
    }

    next->State    = ThreadRunning;
    KiRefillQuantumLocked(next);
    gCurrentThread = next;

    // 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 != gIdleThread && cur->State == ThreadRunning) {
        cur->State = ThreadReady;
        ListInsertTail(&gRunQueue[cur->Priority], &cur->SchedulerLink);
    }

    gOutgoingThread = cur;
    gSchedStats.ContextSwitches++;
    KeSwitchContext(cur, next);
    KeSchedulerPostSwitch();
}

void KiScheduleFromInterrupt(void) {
    if (!KiIsReschedulePending()) return;
    KiSchedule();
}

void KiSchedulerNoteVoluntaryBlock(THREAD* thread) {
    if (!thread) return;
    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    gSchedStats.VoluntaryBlocks++;
    thread->QuantumRemaining = 0;
    KeReleaseSpinlock(&gSchedLock, flags);
}

void KeSchedulerTick(void) {
    THREAD* cur = gCurrentThread;
    if (!cur) return;

    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    cur->SchedLockFlags = flags;

    gSchedStats.Ticks++;

    int requestReschedule = 0;

    if (cur == gIdleThread) {
        if (KiAnyRunnableLocked())
            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 (KiAnyRunnableLocked()) {
                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();
}

__declspec(dllexport) void KeYield(void) {
    KeAssertIrql(PASSIVE_LEVEL);
    KiSchedule();
}