// SPDX-License-Identifier: GPL-3.0-or-later
// Copyright (c) 2026 0xKSor

#include <OS/Scheduler.h>
#include <OS/Panic.h>
#include <Arch/Exceptions.h>
#include <Arch/CPU.h>
#include <VM/VMM.h>
#include <VM/Heap.h>
#include <Lib/Align.h>
#include <Lib/String.h>

OSTask* gOSSchedulerCurrentTask = nullptr;
UInt32 gOSSchedulerNextProcessID = 1;

static OSProcess sOSSchedulerKernelProcess;

void SchedulerInitialize() {
    sOSSchedulerKernelProcess.id = 0;
    sOSSchedulerKernelProcess.state = OSTaskStateRunning;
    sOSSchedulerKernelProcess.l0Table = (Address*)VMPhysToHHDM((Address)gVMKernelL0Table);
    StringCopy(sOSSchedulerKernelProcess.name, "kernel");

    OSTask* kernelTask = (OSTask*)HeapAllocate(sizeof(OSTask));
    if (!kernelTask) OSPanic("Failed to allocate kernel task");

    MemorySet(kernelTask, 0, sizeof(OSTask));
    kernelTask->id = 0;
    kernelTask->process = &sOSSchedulerKernelProcess;
    kernelTask->sleepTicks = 0;
    kernelTask->next = kernelTask;
    kernelTask->state = OSTaskStateRunning;
    kernelTask->waitingForProcessId = -1;

    gOSSchedulerCurrentTask = kernelTask;
}

OSTask* SchedulerSpawn(void(*entryPoint)(), OSProcess* owner, Boolean isUser, Address fixedUserStackAddress) {
    OSTask* task = (OSTask*)HeapAllocate(sizeof(OSTask));
    Pointer stackBase = nullptr;
    if (!task) goto cleanup;

    if (!owner) owner = &sOSSchedulerKernelProcess;
    Size stackSize = kOSSchedulerTaskStackSize;
    stackBase = HeapAllocate(stackSize);
    if (!stackBase) goto cleanup;

    Address stackPointer = (Address)stackBase + stackSize;
    stackPointer -= sizeof(ExceptionsContext);
    stackPointer = AlignDown64(stackPointer, 16);
    ExceptionsContext* context = (ExceptionsContext*)stackPointer;
    MemorySet(context, 0, sizeof(ExceptionsContext));
    context->elr_el1 = (Address)entryPoint;

    if (isUser) {
        context->spsr_el1 = 0x00000000; // El0t (M[3:0] = 0b0000) allow interrupts
        context->sp_el0 = fixedUserStackAddress;
    } else {
        context->spsr_el1 = 0x00000005; // El1h (M[3:0] = 0b0101) allow interrupts
    }

    task->stackPointer = stackPointer;
    task->process = owner;
    task->id = owner->id;
    task->sleepTicks = 0;
    task->kernelStackBase = stackBase;
    task->kernelStackTop = (Address)stackBase + stackSize;
    task->state = OSTaskStateRunning;
    task->waitingForProcessId = -1;

    task->next = gOSSchedulerCurrentTask->next;
    gOSSchedulerCurrentTask->next = task;

    return task;
cleanup:
    if (task) HeapFree(task);
    if (stackBase) HeapFree(stackBase);
    return nullptr;
}

Address SchedulerNext(Address stackPointer) {
    if (!gOSSchedulerCurrentTask) return stackPointer;

    gOSSchedulerCurrentTask->stackPointer = stackPointer;
    OSTask* taskIterator = gOSSchedulerCurrentTask->next;

    do {
        if (taskIterator->sleepTicks > 0) taskIterator->sleepTicks--;
        taskIterator = taskIterator->next;
    } while (taskIterator != gOSSchedulerCurrentTask->next);

    if (gOSSchedulerCurrentTask->sleepTicks > 0) gOSSchedulerCurrentTask->sleepTicks--;

    OSTask* nextTask = gOSSchedulerCurrentTask->next;
    while (true) {
        if (nextTask->state == OSTaskStateDead) {
            nextTask = nextTask->next;
            if (nextTask == gOSSchedulerCurrentTask) OSPanic("No running tasks");
            continue;
        }
        if (nextTask->state == OSTaskStateSleeping && nextTask->sleepTicks == 0) nextTask->state = OSTaskStateRunning;
        if (nextTask->state == OSTaskStateRunning) break;
        nextTask = nextTask->next;

        if (nextTask == gOSSchedulerCurrentTask) {
            if (gOSSchedulerCurrentTask->state == OSTaskStateRunning) break;
            OSPanic("No running tasks");
        }
    }

    if (nextTask->process->l0Table != gOSSchedulerCurrentTask->process->l0Table) {
        Address physicalL0Table = VMHHDMToPhys((Address)nextTask->process->l0Table);
        CPUSwitchAddressSpace(physicalL0Table);
    }
    gOSSchedulerCurrentTask = nextTask;

    return gOSSchedulerCurrentTask->stackPointer;
}

void SchedulerYield(UInt64 ticks) {
    gOSSchedulerCurrentTask->sleepTicks = ticks;
    gOSSchedulerCurrentTask->state = OSTaskStateSleeping;
    CPUException(kOSSchedulerExceptionNumber);
}

UInt64 SchedulerGetNextProcessID() {
    return gOSSchedulerNextProcessID++;
}