// ke_sched.c - MLFQ scheduler
#include <stddef.h>
#include <stdbool.h>
#include <ke_sched.h>
#include <ke_thread.h>
#include <ke_spinlock.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;

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

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

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

    extern KSTATUS KeCreateProcess(const char *, PROCESS **);
    PROCESS* idleProc;
    KeCreateProcess("idle", &idleProc);

    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->SchedLockFlags  = 0x202ULL; // IF=1

    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)IdleLoop;
    idle->SavedContext = ctx;

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

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

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

    cur->QuantaConsumed++;

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

    // 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;
                ListInsertTail(&gRunQueue[PriorityHigh], &t->SchedulerLink);
            }
        }
    }

    THREAD* next = PickNext();
    if (next == cur) {
        if (cur != gIdleThread) {
            ListInsertTail(&gRunQueue[cur->Priority], &cur->SchedulerLink);
        }
        KeReleaseSpinlock(&gSchedLock, flags);
        return;
    }

    cur->State  = ThreadReady;
    next->State = ThreadRunning;
    gCurrentThread = next;

    if (cur != gIdleThread && cur->State == ThreadReady) {
        ListInsertTail(&gRunQueue[cur->Priority], &cur->SchedulerLink);
    }

    gOutgoingThread = cur;
    KeSwitchContext(cur, next);
    KeSchedulerPostSwitch();
}

__declspec(dllexport) void KeYield(void) {
    THREAD* cur = gCurrentThread;
    if (!cur) return;

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

    THREAD* next = PickNext();
    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;
    KeSwitchContext(cur, next);
    KeSchedulerPostSwitch();
}