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

#define DEMOTION_QUANTA  4 /* ticks at a level before demotion */

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;

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->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);
    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()) {
        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;
    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 KeSchedulerTick(void) {
    THREAD* cur = gCurrentThread;
    if (!cur) return;

    uintptr_t flags = KeAcquireSpinlock(&gSchedLock);
    cur->SchedLockFlags = flags; // Issue 5 fix

    gSchedStats.Ticks++;
    cur->QuantaConsumed++;

    if (cur != gIdleThread && cur->QuantaConsumed >= DEMOTION_QUANTA) {
        cur->QuantaConsumed = 0;
        if (cur->Priority > PriorityLow) {
            cur->Priority = (THREAD_PRIORITY)(cur->Priority - 1);
            gSchedStats.Demotions++;
        }
    }

    // Periodic aging pass to prevent indefinite starvation of PriorityLow threads
    gAgingCounter++;
    if (gAgingCounter >= 100) {
        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;
                gSchedStats.AgingPromotions++;
                ListInsertTail(&gRunQueue[PriorityHigh], &t->SchedulerLink);
            }
        }
    }

    KeReleaseSpinlock(&gSchedLock, flags);

    KiRequestReschedule();
}

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