MontaukOS/kernel/src/Sched/Scheduler.hpp

/*
    * Scheduler.hpp
    * Preemptive process scheduler with SMP support
    * Copyright (c) 2025-2026 Daniel Hammer
*/

#pragma once
#include <cstdint>
#include <Api/Syscall.hpp>

namespace Sched {

    static constexpr int MaxProcesses = 256;
    static constexpr uint64_t StackPages = 4;  // 16 KiB kernel stack per process
    static constexpr uint64_t StackSize = StackPages * 0x1000;
    static constexpr uint64_t UserStackPages = 8;  // 32 KiB user stack
    static constexpr uint64_t UserStackSize = UserStackPages * 0x1000;
    static constexpr uint64_t UserStackTop = 0x7FFFFFF000ULL;  // User stack top VA
    static constexpr uint64_t UserHeapBase = 0x40000000ULL;    // User heap start VA
    static constexpr uint32_t UserReadDirSlots = 64;           // rotating scratch pages for SYS_READDIR
    static constexpr uint64_t UserReadDirBase =
        UserHeapBase - (uint64_t)UserReadDirSlots * 0x1000ULL;
    static constexpr uint64_t ExitStubAddr = 0x3FF000ULL;     // User-space exit stub page
    static constexpr uint64_t TimeSliceMs = 10; // 10 ms time slice

    enum class ProcessState {
        Free,
        Ready,
        Running,
        Blocked,       // Waiting (e.g. waitpid) -- not schedulable
        Terminated
    };

    struct Process {
        int pid;
        ProcessState state;
        int waitingForPid;     // PID this process is blocked on (-1 if none)
        uint64_t sleepUntilTick; // Tick deadline for sleep/object wait timeout (0 = none)
        void* waitingOnObject; // IPC/scheduler object this process is blocked on (nullptr if none)
        char name[64];
        uint64_t savedRsp;
        uint64_t stackBase;       // Bottom of allocated kernel stack (lowest address)
        uint64_t entryPoint;
        uint64_t sliceRemaining;  // Milliseconds left in current time slice
        uint64_t pml4Phys;        // Physical address of per-process PML4
        uint64_t kernelStackTop;  // Top of kernel stack (for TSS RSP0 / SYSCALL)
        uint64_t userStackTop;    // User-space stack top
        uint64_t heapNext;        // Simple bump allocator for user heap
        uint32_t readdirCursor;   // Next SYS_READDIR scratch slot
        char args[256];           // Command-line arguments (set by parent via Spawn)
        char user[32];            // Owner user name (inherited from parent on spawn)
        char cwd[256];            // Absolute current working directory

        int runningOnCpu;         // CPU index running this process (-1 if not running)
        bool killPending = false; // Set by Sys_Kill when target is running on another CPU
        bool reapReady = false;   // Set once teardown is complete and BSP may free slot resources

        // I/O redirection for GUI terminal
        bool redirected = false;
        int parentPid = -1;
        uint8_t* outBuf = nullptr;   // 4KB ring: child writes (print/putchar), parent reads
        volatile uint32_t outHead = 0;
        volatile uint32_t outTail = 0;
        uint8_t* inBuf = nullptr;    // 4KB ring: parent writes, child reads (getchar)
        volatile uint32_t inHead = 0;
        volatile uint32_t inTail = 0;
        Montauk::KeyEvent keyBuf[64]; // parent injects, child reads (getkey/iskeyavailable)
        volatile uint32_t keyHead = 0;
        volatile uint32_t keyTail = 0;
        static constexpr uint32_t IoBufSize = 4096;

        // GUI terminal dimensions (set by desktop, read by SYS_TERMSIZE)
        int termCols = 0;
        int termRows = 0;

        // IPC-backed redirected terminal channels
        int ioOutHandle = -1;
        int ioInHandle = -1;
        int ioKeyHandle = -1;
        int ioWaitsetHandle = -1;

        // FPU/SSE state (FXSAVE format, must be 16-byte aligned)
        uint8_t fpuState[512] __attribute__((aligned(16)));
    };

    void Initialize();
    int Spawn(const char* vfsPath, const char* args = nullptr);
    void Schedule();

    // True when there is runnable work somewhere in the process table.
    bool HasReadyProcesses();

    // Called from the APIC timer handler with the elapsed time for that CPU's
    // tick interval. The BSP runs at 1 ms; APs may use a coarser interval.
    void Tick(uint32_t elapsedMs = 1);

    // Get the PID of the currently running process (-1 if idle)
    int GetCurrentPid();

    // Get a pointer to the currently running process (nullptr if idle)
    Process* GetCurrentProcessPtr();

    // Called by terminated processes to mark themselves done
    void ExitProcess();

    // Check if a process is still alive (Ready, Running, or Blocked)
    bool IsAlive(int pid);

    // Block the current process until the given PID exits.
    void BlockOnPid(int pid);

    // Block the current process for the given number of milliseconds.
    void BlockForSleep(uint64_t ms);

    // Block the current process until the given object is signaled or timed out.
    // timeoutMs == 0 means wait indefinitely.
    void BlockOnObject(void* object, uint64_t timeoutMs = 0);

    // Atomically check a condition under the scheduler lock and block only if
    // it still holds. This prevents lost wakeups for edge-triggered waiters.
    bool BlockOnObjectIf(void* object, uint64_t timeoutMs,
                         bool (*shouldBlock)(void*), void* context);

    // BSP-only scheduler housekeeping: wake expired sleepers and reclaim
    // terminated process resources.
    void RunBspMaintenance();

    // Return the earliest blocked sleep/object timeout deadline in ticks,
    // or 0 when no timed waits are pending.
    uint64_t GetNextDeadlineTick();

    // Wake any processes blocked on the given object.
    void WakeObjectWaiters(void* object);

    // Kill a process by PID. If the process is running on another CPU,
    // sets a kill-pending flag checked on the next timer tick.
    // Returns 0 on success, -1 on failure.
    int KillProcess(int pid);

    // Find a process by PID (returns nullptr if not found or not alive)
    Process* GetProcessByPid(int pid);

    // Spawn a crashpad process with the given crash file path as argument.
    // Called by the exception handler when a process faults.
    int SpawnCrashPad();

    // Get a pointer to slot i in the process table (for enumeration)
    Process* GetProcessSlot(int slot);

    // Per-process allocated page count (tracked by Heap syscalls, separate from Process struct)
    inline uint64_t g_allocatedPages[MaxProcesses] = {};

}