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

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

namespace Sched {

    // Process and thread slots share the same table: the main thread of a
    // process owns the per-process resources (PML4, IPC handles, args/cwd/user,
    // redirected I/O), while sibling threads occupy additional slots whose
    // `primarySlot` points back to the main thread. The scheduler treats
    // every slot uniformly -- the primary/sibling distinction matters only
    // for resource ownership on exit.
    static constexpr int MaxProcesses = 256;
    // 64 KiB kernel stack per thread. The AML interpreter is the driver of
    // this size: it runs on the calling thread's stack when a syscall asks for
    // battery status, and firmware methods nest deeply through frames that
    // carry 256-byte path buffers and 272-byte AML objects. See
    // AML::StackBudget, which must stay comfortably under this.
    static constexpr uint64_t StackPages = 16;
    static constexpr uint64_t StackSize = StackPages * 0x1000;
    static constexpr uint64_t UserStackPages = 8;  // 32 KiB user stack (main thread, eagerly mapped)
    static constexpr uint64_t UserStackSize = UserStackPages * 0x1000;
    static constexpr uint64_t UserStackTop = 0x7FFFFFF000ULL;  // Main-thread user stack top VA
    // The main-thread stack demand-grows below the eager region up to this
    // limit: page faults in [UserStackTop - UserStackMax, UserStackTop -
    // UserStackSize) map fresh zeroed pages instead of killing the process.
    // Heavy recursive programs (GCC's cc1plus) need megabytes of stack.
    static constexpr uint64_t UserStackMax = 64 * 1024 * 1024;  // 64 MiB
    static constexpr uint64_t UserHeapBase = 0x40000000ULL;    // User heap start VA
    // Keep the bump-allocated heap below the fixed shared-library range.  A
    // failed bound check is recoverable (SYS_ALLOC returns null); allowing the
    // heap to cross this boundary would silently replace library PTEs.
    static constexpr uint64_t UserHeapLimit = 0x60000000ULL;
    static constexpr uint64_t UserHeapPages =
        (UserHeapLimit - UserHeapBase) / 0x1000ULL;
    static constexpr uint64_t UserHeapBitmapWords = (UserHeapPages + 63) / 64;
    // Surface mappings use reusable, fixed-size per-process slots instead of
    // consuming heap VA forever.  64 * 32 MiB occupies [0x7000000000,
    // 0x7080000000), comfortably below the main user stack.
    static constexpr uint64_t UserSurfaceBase = 0x7000000000ULL;
    static constexpr uint64_t UserSurfaceSlotSize = 32ULL * 1024 * 1024;
    static constexpr uint64_t UserSurfaceSlots = 64;
    static constexpr uint64_t UserSurfaceEnd =
        UserSurfaceBase + UserSurfaceSlots * UserSurfaceSlotSize;
    // The two write-combined scanout mappings live in their own region.
    static constexpr uint64_t UserFramebufferBase = 0x7100000000ULL;
    static constexpr uint64_t UserFramebufferLimit = 0x7140000000ULL;
    static_assert(UserHeapBase < UserHeapLimit);
    static_assert(UserSurfaceEnd <= UserFramebufferBase);
    static_assert(UserFramebufferLimit < UserStackTop - UserStackMax);
    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 {
        static constexpr uint32_t EnvironmentBytes = 4096;
        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 cpuTimeMs;       // Accumulated scheduler runtime in milliseconds
        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;        // High-water mark of the user-heap address space
        uint32_t readdirCursor;   // Next SYS_READDIR scratch slot
        char args[4096];          // Command-line arguments (set by parent via Spawn)
        char environment[EnvironmentBytes]; // NUL-separated NAME=VALUE entries
        uint32_t environmentLength;
        char user[32];            // Owner user name (inherited from parent on spawn)
        uint64_t effectiveCaps;   // Authority currently usable by syscalls
        uint64_t delegableCaps;   // Authority this process may pass to children
        int sessionId;            // Process-session leader PID (inherited on spawn)
        char cwd[256];            // Absolute current working directory

        // First-use consent (Sched/Consent.cpp). Primary slot only.
        uint64_t consentCeiling;  // eligible caps the parent could have delegated at spawn
        uint64_t consentDenied;   // caps refused for the rest of this process's life
        char imagePath[256];      // full resolved image path; name[] truncates

        // Thread-local storage. fsBase is loaded into IA32_FS_BASE when
        // this thread is dispatched (0 = no TLS). The tls* template fields
        // live on the primary slot and describe PT_TLS so sibling threads
        // can get their own block copies.
        uint64_t fsBase = 0;
        uint64_t tlsTemplateVaddr = 0;
        uint64_t tlsFileSize = 0;
        uint64_t tlsMemSize = 0;
        uint64_t tlsAlign = 0;

        int runningOnCpu;         // CPU index running this process (-1 if not running)
        std::atomic<bool> killPending{false}; // Cross-CPU kill handoff
        bool reapReady = false;   // Set once teardown is complete and BSP may free slot resources
        bool startPending = false; // Spawned but not yet made runnable

        // Threading. Every slot is a schedulable thread. The "main thread" of a
        // process has primarySlot == own slot and owns per-process resources
        // (pml4Phys, IPC handles, libs, heap, redirected I/O, cwd/user/args).
        // Sibling threads share those resources via primarySlot.
        int primarySlot = -1;       // Index into processTable; -1 if Free
        int joinerSlot = -1;        // Slot blocked in JoinThread on this slot
        int exitCode = 0;           // Returned to joiner / accumulated for the process
        bool joinable = true;       // Sibling thread keeps Terminated state for the joiner
        uint64_t threadArg = 0;     // Argument passed into entryPoint on first dispatch

        // Futex wait state (see FutexWait). futexAddr is the user VA this
        // thread is parked on, 0 when not in a futex wait. futexSeq orders
        // waiters so FutexWake releases the longest-waiting first.
        uint64_t futexAddr = 0;
        uint64_t futexSeq = 0;
        bool futexWoken = false;

        // 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;
        volatile uint32_t keyHead = 0;
        volatile uint32_t keyTail = 0;
        // SYS_GETCHAR returns a key's UTF-8 encoding one byte per call; the
        // bytes not handed out yet wait here.
        uint8_t getcharUtf8[4] = {};
        uint8_t getcharUtf8Len = 0;
        uint8_t getcharUtf8Pos = 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, bool startReady = true,
              const char* environment = nullptr, uint32_t environmentLength = 0,
              const montauk::abi::SpawnCapabilities* capabilities = nullptr,
              const char* userOverride = nullptr);
    int StartProcess(int pid);
    void Schedule();

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

    // Ask the current CPU to yield once to its idle kernel context on the next
    // scheduler pass so deferred device work cannot starve under full load.
    void RequestKernelWork();

    // 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, bool interruptedUser = false);

    // Get the PID of the currently running process (-1 if idle).
    // For sibling threads this returns the primary slot's PID -- i.e. the
    // process-level identifier, which is what userspace sees.
    int GetCurrentPid();

    // Get the per-thread id of the currently running thread (-1 if idle).
    // For the main thread this equals GetCurrentPid().
    int GetCurrentTid();

    // Get a pointer to the currently running process (primary slot, nullptr if idle).
    // Always returns the slot that owns per-process state -- never a sibling thread.
    Process* GetCurrentProcessPtr();

    // Capability checks always consult kernel-owned process metadata. User
    // names are deliberately excluded from authorization.
    bool HasCapability(uint64_t capability);

    // Get a pointer to the currently running thread's slot (may be a sibling).
    Process* GetCurrentThreadPtr();

    // Demand-grow the current process's main-thread user stack. Called from
    // the page fault handler when a non-present fault lands in the stack
    // growth region. Maps one zeroed page at the faulting address and returns
    // true if the faulting instruction should be retried.
    bool TryGrowUserStack(uint64_t faultAddr);

    // Called by terminated processes to mark themselves done.
    // Tears down the entire process (kills all sibling threads, frees address space).
    void ExitProcess();

    // ====================================================================
    // Threading
    // ====================================================================

    // Spawn a sibling thread inside the current process.
    // `entry` is a user-mode RIP, `arg` is delivered to entry in %rdi, and
    // `userStackTop` points to the top-of-stack the caller (libc) allocated
    // out of the user heap. Returns the new TID, or -1 on failure.
    int SpawnThread(uint64_t entry, uint64_t arg, uint64_t userStackTop);

    // Reserve a page-aligned range from a process's shared bump heap.  The
    // reservation is serialized across sibling threads and bounded below the
    // fixed shared-library mappings.
    bool ReserveUserHeapRange(int primarySlot, uint64_t size, uint64_t& outVa);
    void ReleaseUserHeapRange(int primarySlot, uint64_t va, uint64_t size);

    // Terminate the currently executing thread. If this is the main thread,
    // the entire process exits (equivalent to ExitProcess).
    [[noreturn]] void ExitCurrentThread(int exitCode);

    // Block the current thread until `tid` (a sibling of the same process)
    // terminates. On success, writes the exit code via outExitCode (may be
    // null) and frees the joined thread's slot. Returns 0 on success, -1 if
    // tid is invalid or not joinable by the caller.
    int JoinThread(int tid, int* outExitCode);

    // 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);
    int LookupExitCode(int pid);
    void SetProcessExitCode(int pid, int code);

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

    // Take this token before testing an object's ready condition, then pass it
    // to BlockOnObjectSince. A wake between the test and scheduler enrollment
    // makes the latter return without blocking instead of losing the edge.
    uint64_t ObserveObjectWake(void* object);
    bool BlockOnObjectSince(void* object, uint64_t timeoutMs, uint64_t observedWake);

    // BSP idle-context housekeeping: reclaim terminated process resources and
    // run policy work previously queued by the timer or process teardown.
    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);

    // ====================================================================
    // Futexes
    // ====================================================================

    // Futexes are private to a process: the key is (primary slot, user VA),
    // so two processes waiting on the same address never see each other.
    static constexpr uint64_t FutexWaitForever = ~0ULL;

    // Block the current thread on the 32-bit word at `uaddr` if it still
    // holds `expected`. The comparison and the enrollment happen atomically
    // under the scheduler lock, so a FutexWake issued after the waker changed
    // the word can never be lost. `timeoutNs` is rounded up to the 1 ms tick;
    // FutexWaitForever waits without a deadline, 0 only performs the check.
    // Returns Success when woken by FutexWake, WouldBlock when the word did
    // not hold `expected`, TimedOut when the deadline passed, InvalidArgument
    // for a misaligned address, BadPointer when the word is not mapped.
    // Spurious Success returns are not possible, but callers must still
    // re-check their condition: another thread may have run in between.
    int FutexWait(uint64_t uaddr, uint32_t expected, uint64_t timeoutNs);

    // Wake up to `count` threads of the current process blocked on `uaddr`,
    // longest-waiting first. Returns the number woken.
    int FutexWake(uint64_t uaddr, uint32_t count);

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

    // Start a new process session for the caller, or terminate every process
    // belonging to a session. KillSession returns the number of live members
    // it signalled; callers can repeat until it returns zero.
    int CreateSession();
    int KillSession(int sessionId);

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

    // Consent support for Sched/Consent.cpp. Each validates the (subject,
    // agent) pair and reads or mutates both under the scheduler lock, so a
    // grant can never land on a slot that exited and was reused.
    //
    // The pair is valid when both are live, distinct, in the same session,
    // the agent still holds CAP_COMPOSITOR and has `capability` delegable,
    // and `capability` lies within the subject's consent ceiling.
    bool ConsentSubjectEligible(int pid, int agentPid, uint64_t capability);
    bool ApplyConsentGrant(int pid, int agentPid, uint64_t capability);
    void ApplyConsentDenial(int pid, uint64_t capability);
    // Fill imagePath/imageProtected for the request; false if pid is gone.
    bool SnapshotConsentSubject(int pid, montauk::abi::ConsentRequest& out);

    // Spawn the crashpad to report on a faulted process (pid `crasherPid`).
    // Called by the exception handler when a process faults. If the crasher had
    // a redirected console, the crashpad inherits its out-stream so it can print
    // a textual report to a console session with no desktop.
    int SpawnCrashPad(int crasherPid);

    // 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] = {};
    // One bit per page in the bounded userspace heap. This is the authoritative
    // allocation state; unlike heapNext, it makes virtual ranges reusable after
    // unmap and failed reservations.
    inline uint64_t g_userHeapPageMap[MaxProcesses][UserHeapBitmapWords] = {};

}