#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>

#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
    void* Teb;

    // 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];
} 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 idle 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.
void    KeBlockThread(THREAD* thread, OBJECT* waitObject);
void    KeUnblockThread(THREAD* thread);

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

void    KeInitThreadFpuContext(THREAD* thread);