#pragma once
// ke_thread.h - thread model
// A THREAD is the unit of scheduling. Every thread belongs to exactly one process.
#include <stdint.h>
#include <ob_object.h>
#include <ke_list.h>
#include <ke_process.h>
#include <ke_event.h>
#include <ke_spinlock.h>
#include <ke_trap.h>

#define KSTACK_SIZE   (64  * 1024)    /* 64 KB kernel stack  */
#define USTACK_SIZE   (512 * 1024)    /* 512 KB user stack   */

typedef enum THREAD_PRIORITY {
    PriorityLow    = 0,
    PriorityNormal = 1,
    PriorityHigh   = 2,
    PriorityCount  = 3,
} THREAD_PRIORITY;

typedef enum THREAD_STATE {
    ThreadReady   = 0,
    ThreadRunning = 1,
    ThreadWaiting = 2,
    ThreadDead    = 3,
} THREAD_STATE;

// Saved register context - matches the layout pushed by KeSwitchContext.
// The saved RSP points to this structure on the kernel stack.
typedef struct KCONTEXT {
    uint64_t R15, R14, R13, R12;
    uint64_t Rbp, Rbx;
    uint64_t Rip;      /* return address back into KeSwitchContext */
} KCONTEXT;

typedef struct THREAD {
    OBJECT          Header;            /* must be first */

    uint32_t        ThreadId;
    THREAD_STATE    State;
    THREAD_PRIORITY Priority;

    // Quanta consumed without blocking - used by the MLFQ for demotion.
    uint32_t        QuantaConsumed;

    PROCESS*        Owner;
    LIST_ENTRY      ProcessLink;       /* node in Owner->ThreadList  */
    LIST_ENTRY      SchedulerLink;     /* node in run queue or wait list */

    // Kernel stack - never PAGE_USER, pointed to by TssRsp0.
    void*           KernelStackBase;
    uintptr_t       KernelStackTop;

    // Saved context pointer - set by KeSwitchContext.
    KCONTEXT* SavedContext;

    // Non-NULL while State == ThreadWaiting.
    OBJECT* WaitObject;

    // Updated on every context switch so the CPU loads the right kernel stack.
    uintptr_t       TssRsp0;

    // Thread Environment Block. GS base is loaded from this pointer.
    void*           Teb;
    size_t          TebSize;

    // Ring-3 entry point and argument for newly created threads.
    uintptr_t       UserEntry;
    uintptr_t       UserArg;

    uintptr_t       SchedLockFlags; // Stores RFLAGS while thread is context switched out
    void*           UserStackBase;  // Tracks user stack allocation to prevent memory leaks

    // Signaled (manual-reset) when the thread exits.
    // Allocated in KeCreateThread, signaled in KeExitThread before teardown.
    // Waiters block on this via SYS_WAIT_THREAD -> KeWaitForEvent.
    EVENT*          ExitEvent;

    void*   FpuContext;
    uint8_t FpuStorage[528];

    // Set by wait primitives before blocking and completed by the object
    // that wakes the thread.  Keep this field after the assembly-visible
    // context-switch fields: ke_switch.asm hard-codes offsets through
    // FpuContext.
    KSTATUS WaitStatus;
    LIST_ENTRY* WaitListHead;
    uint8_t AlertableWait;

    // Kernel APC state.  APCs are queued to a target thread and delivered
    // when that thread reaches PASSIVE_LEVEL with APC delivery enabled.
    LIST_ENTRY ApcQueue;
    SPINLOCK   ApcLock;
    volatile long KernelApcPending;
    int32_t    KernelApcDisable;
    uint8_t    ApcDeliveryActive;

    // Non-NULL while executing on a syscall/interrupt/fault entry frame.
    // The pointed-to frame lives on this thread's kernel stack.
    KTRAP_FRAME* TrapFrame;

    // User APC state.  User APCs are queued by kernel completions, but are
    // delivered only on an alertable return to user mode through the registered
    // user-mode dispatcher stub.
    LIST_ENTRY UserApcQueue;
    SPINLOCK   UserApcLock;
    volatile long UserApcPending;
    uintptr_t  UserApcDispatcher;
    uint8_t    UserApcDeliverable;
    uint8_t    UserApcDeliveryActive;
} THREAD;

__declspec(dllexport) KSTATUS KeCreateThread(PROCESS* owner, uintptr_t entry, uintptr_t arg,
                       THREAD_PRIORITY priority, THREAD** out);

// Create a kernel-mode-only thread. fn is called directly in ring 0 with
// arg; when fn returns the thread exits automatically. No user stack or TEB
// is allocated. The thread is owned by the system process and runs at the
// given priority.
__declspec(dllexport) KSTATUS KeCreateKernelThread(uintptr_t fn, uintptr_t arg,
                              THREAD_PRIORITY priority, THREAD** out);

// Mark a thread dead and deschedule it permanently.
__declspec(dllexport) void    KeExitThread(THREAD* thread, int code);

// Block / unblock primitives used by sync objects.
//
// Correct wait pattern:
//   object lock held
//   KePrepareBlockThread(self, object, alertable);
//   insert self into object's wait list
//   release object lock
//   KeCommitBlockThread(self, alertable);
//
// The split form prevents the classic lost wakeup where an ISR/DPC signals
// the object after the waiter is queued but before the waiter marks itself
// waiting.  Alertable waits are explicit at every call site.
void    KePrepareBlockThread(THREAD* thread, OBJECT* waitObject, uint8_t alertable);
void    KeSetThreadWaitList(THREAD* thread, LIST_ENTRY* waitListHead);
KSTATUS KeCommitBlockThread(THREAD* thread, uint8_t alertable);
KSTATUS KeBlockThread(THREAD* thread, OBJECT* waitObject, uint8_t alertable);
void    KeUnblockThreadWithStatus(THREAD* thread, KSTATUS status);
void    KeUnblockThread(THREAD* thread);
KSTATUS KeAlertThread(THREAD* thread);
void    KiWakeThreadForKernelApc(THREAD* thread);
void    KiWakeThreadForUserApc(THREAD* thread);

// Assembly stub - implemented in ke_switch.asm.
void    KeSwitchContext(THREAD* from, THREAD* to);

void    KeInitThreadFpuContext(THREAD* thread);